llama : do not map foreign ranges into a backend buffer (#30405)

* llama : do not map foreign ranges into a backend buffer

A buffer created from a host pointer pins everything that it wraps. When
several contexts share the same file (lazy tensors, MoE tensors kept on the
CPU with -cmoe), the single buffer that spans the first and the last tensor
of a context also wraps the ranges owned by the other contexts, pinning data
that the backend never reads.

For models with per-layer token embeddings this pins the whole lazy tensor
into the GPU buffer (27 GiB for Qwen3.8-Flash-Next IQ4_NL), defeating the
point of reading it lazily.

Instead, compute the file ranges of each context and subtract the ranges
owned by the other contexts, then create one buffer per range. Weights that
are mapped in more than one context (e.g. token_embd.weight) keep being part
of the range of every context that uses them.

Assisted-by: pi:llama.cpp/Qwen3.8-Flash-Next

* llama : enable prefetch on macOS

The prefetch hint of the tensors marked as prefetchable is a no-op on macOS,
so their pages are faulted in one by one while the graph runs. madvise
(MADV_WILLNEED) behaves as expected here, so add the platform to the guard.

With the lazy per-layer token embeddings of Qwen3.8-Flash-Next IQ4_NL read
from a cold file cache, prompt processing goes from 356 / 375 / 475 t/s to
735 / 749 / 788 t/s for pp512 / pp2048 / pp4096.

Assisted-by: pi:llama.cpp/Qwen3.8-Flash-Next

* llama : use a plain buffer list per context

Assisted-by: pi:llama.cpp/Qwen3.8-Flash-Next

.github/workflows/build-apple.yml

@@ -105,7 +105,7 @@ jobs:
           cd build
           # Metal Paravirtual devices are difficult to support -> disable
           # ref: https://github.com/ggml-org/llama.cpp/pull/19802#issuecomment-4013704023
-          ctest -L main -E "test-llama-archs|test-save-load-state|test-recurrent-state-rollback" --verbose --timeout 900
+          ctest -L main -E "test-llama-archs|test-llama-context" --verbose --timeout 900
 
   macos-latest-x64:
     runs-on: macos-15-intel

.github/workflows/build-openvino.yml

@@ -37,7 +37,7 @@ env:
   LLAMA_ARG_LOG_PREFIX: 1
   LLAMA_ARG_LOG_TIMESTAMPS: 1
   # TODO: fix failing tests on OpenVINO backend
-  CTEST_EXCLUDE: "test-llama-archs|^test-recurrent-state-|test-save-load-state"
+  CTEST_EXCLUDE: "test-llama-archs|test-llama-context"
 
 jobs:
   ubuntu-24-openvino:

.github/workflows/release.yml

@@ -714,6 +714,9 @@ jobs:
             -DGGML_OPENVINO=ON \
             -DCMAKE_INSTALL_RPATH='$ORIGIN' \
             -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
+            -DGGML_BACKEND_DL=ON \
+            -DGGML_NATIVE=OFF \
+            -DGGML_CPU_ALL_VARIANTS=ON \
             ${{ env.CMAKE_ARGS }}
           cmake --build build/ReleaseOV --config Release --parallel
 
@@ -752,8 +755,10 @@ jobs:
         with:
           key: release-ubuntu-24.04-openvino-release-no-preset-v1
 
+  # note: builds only the ggml-openvino backend - llama-server and the CPU backend are injected from the
+  #       windows-cpu zip during the release "Merge artifacts" step
   windows-openvino:
-    needs: [check-release, ui-build]
+    needs: [check-release]
     if: ${{ needs.check-release.outputs.should_release == 'true' }}
 
     runs-on: windows-2022
@@ -770,12 +775,6 @@ jobs:
         with:
             fetch-depth: 0
 
-      - name: Download UI build
-        uses: actions/download-artifact@v8
-        with:
-          name: llama-ui.zip
-          path: tools/ui/dist
-
       - name: ccache
         uses: ggml-org/ccache-action@v1.2.24
         with:
@@ -823,11 +822,13 @@ jobs:
             -A x64 ^
             -DCMAKE_BUILD_TYPE=Release ^
             -DGGML_OPENVINO=ON ^
+            -DGGML_BACKEND_DL=ON ^
+            -DGGML_NATIVE=OFF ^
+            -DGGML_CPU=OFF ^
             -DLLAMA_BUILD_BORINGSSL=ON ^
-            -DCMAKE_TOOLCHAIN_FILE=C:\vcpkg\scripts\buildsystems\vcpkg.cmake ^
-            ${{ env.CMAKE_ARGS }}
+            -DCMAKE_TOOLCHAIN_FILE=C:\vcpkg\scripts\buildsystems\vcpkg.cmake
 
-          cmake --build build\ReleaseOV --config Release -- /m
+          cmake --build build\ReleaseOV --config Release --target ggml-openvino -- /m
 
       - name: Pack artifacts
         id: pack_artifacts
@@ -855,12 +856,12 @@ jobs:
           Copy-Item -Path (Join-Path $OPENVINO_ROOT 'docs\licensing\*') -Destination $licensingDest -Recurse -Force
 
           Copy-Item LICENSE $dest
-          7z a -snl llama-${{ needs.check-release.outputs.tag_name }}-bin-win-openvino-${{ env.OPENVINO_VERSION_MAJOR }}-x64.zip $dest\*
+          7z a -snl llama-bin-win-openvino-${{ env.OPENVINO_VERSION_MAJOR }}-x64.zip $dest\*
 
       - name: Upload artifacts
         uses: actions/upload-artifact@v7
         with:
-          path: llama-${{ needs.check-release.outputs.tag_name }}-bin-win-openvino-${{ env.OPENVINO_VERSION_MAJOR }}-x64.zip
+          path: llama-bin-win-openvino-${{ env.OPENVINO_VERSION_MAJOR }}-x64.zip
           archive: false
 
       - name: ccache-clear

ci/run.sh

@@ -180,7 +180,7 @@ if [ ! -z ${GG_BUILD_OPENVINO} ]; then
     CMAKE_EXTRA="${CMAKE_EXTRA} -DGGML_OPENVINO=ON"
 
     # TODO: fix failing tests on OpenVINO backend
-    CTEST_EXTRA="-E test-llama-archs|^test-recurrent-state-|test-save-load-state"
+    CTEST_EXTRA="-E test-llama-archs|test-llama-context"
 fi
 
 ## helpers
@@ -447,10 +447,10 @@ function gg_run_qwen3_0_6b {
 
     (time ./bin/llama-imatrix --model ${model_f16} -f ${wiki_test} -ngl 99 -c 1024 -b 512 --chunks 2 ) 2>&1 | tee -a $OUT/${ci}-imatrix.log
 
-    (time ./bin/test-save-load-state --model ${model_q4_0} -ngl 10 -c 1024 -fa off --no-op-offload) 2>&1 | tee -a $OUT/${ci}-save-load-state.log
-    (time ./bin/test-save-load-state --model ${model_q4_0} -ngl 10 -c 1024 -fa on  --no-op-offload) 2>&1 | tee -a $OUT/${ci}-save-load-state.log
-    (time ./bin/test-save-load-state --model ${model_q4_0} -ngl 99 -c 1024 -fa off                ) 2>&1 | tee -a $OUT/${ci}-save-load-state.log
-    (time ./bin/test-save-load-state --model ${model_q4_0} -ngl 99 -c 1024 -fa on                 ) 2>&1 | tee -a $OUT/${ci}-save-load-state.log
+    (time ./bin/test-llama-context --model ${model_q4_0} -ngl 10 -c 1024 -fa off --no-op-offload) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
+    (time ./bin/test-llama-context --model ${model_q4_0} -ngl 10 -c 1024 -fa on  --no-op-offload) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
+    (time ./bin/test-llama-context --model ${model_q4_0} -ngl 99 -c 1024 -fa off                ) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
+    (time ./bin/test-llama-context --model ${model_q4_0} -ngl 99 -c 1024 -fa on                 ) 2>&1 | tee -a $OUT/${ci}-test-llama-context.log
 
     function check_ppl {
         qnt="$1"

common/log.cpp

@@ -162,6 +162,7 @@ struct common_log {
     // default capacity
     common_log(size_t capacity = 512) {
         file       = nullptr;
+        stdo       = true;
         colors     = false;
         prefix     = false;
         timestamps = false;
@@ -190,6 +191,7 @@ private:
 
     FILE * file;
 
+    bool stdo;
     bool colors;
     bool prefix;
     bool timestamps;
@@ -205,7 +207,9 @@ private:
     bool print_entry(const common_log_entry & e) const {
         if (e.is_end) return true;
 
-        e.print();
+        if (stdo) {
+            e.print();
+        }
         if (file) {
             e.print(file);
         }
@@ -400,6 +404,14 @@ public:
         resume();
     }
 
+    void set_stdo(bool enabled) {
+        pause();
+
+        stdo = enabled;
+
+        resume();
+    }
+
     bool get_colors() const {
         return colors;
     }
@@ -497,6 +509,10 @@ void common_log_set_file(struct common_log * log, const char * file) {
     log->set_file(file);
 }
 
+void common_log_set_stdo(struct common_log * log, bool enabled) {
+    log->set_stdo(enabled);
+}
+
 void common_log_set_colors(struct common_log * log, log_colors colors) {
     if (colors == LOG_COLORS_AUTO) {
         log->set_colors(tty_can_use_colors());

common/log.h

@@ -92,6 +92,7 @@ void common_log_add(struct common_log * log, enum ggml_log_level level, const ch
 //
 
 void common_log_set_file      (struct common_log * log, const char * file); // not thread-safe
+void common_log_set_stdo      (struct common_log * log, bool enabled);      // whether to output to stdout/stderr
 void common_log_set_colors    (struct common_log * log, log_colors colors); // not thread-safe
 bool common_log_get_colors    (struct common_log * log);                    // whether colors are enabled
 void common_log_set_prefix    (struct common_log * log, bool prefix);       // whether to output prefix to each log

convert_lora_to_gguf.py

@@ -127,6 +127,9 @@ class LoraTorchTensor:
         assert dim is None
         return self.shape
 
+    def dim(self) -> int:
+        return len(self.shape)
+
     def contiguous(self) -> LoraTorchTensor:
         return LoraTorchTensor(
             self._lora_A.contiguous(),

docs/backend/zDNN.md

@@ -27,6 +27,9 @@ The llama.cpp zDNN backend is designed to enable llama.cpp on IBM z17 and later
 | F32       | Supported |
 | F16       | Supported |
 | BF16      | Supported |
+| Q8_0      | Supported |
+
+Note: Quantized weights are requantized to int8 with one scale per group of up to 2048 values in each row (or 1024 with `GGML_ZDNN_HIGH_PREC=1`), after each input channel is divided by its RMS (the activations are multiplied by it instead), and each activation is quantized to two int8 values (about 15 bits). This requires `NNPA_PARMBLKFORMAT_1` (IBM z17 / LinuxONE 5).
 
 ## CMake Options
 
@@ -37,6 +40,12 @@ The IBM zDNN backend has the following CMake options that control the behaviour
 | `GGML_ZDNN`  | `OFF`         | Compile llama.cpp with zDNN support |
 | `ZDNN_ROOT`  | `""`          | Override zDNN library lookup        |
 
+## Environment Variables
+
+| Environment Variable  | Default Value | Description                                                                            |
+| --------------------- | ------------- | -------------------------------------------------------------------------------------- |
+| `GGML_ZDNN_HIGH_PREC` | `0`           | Set to `1` to requantize Q8_0 weights in smaller groups, for lower loss at lower speed |
+
 ## 1. Install zDNN Library
 
 Note: Using the zDNN library provided via `apt` or `yum` may not work correctly as reported in [#15772](https://github.com/ggml-org/llama.cpp/issues/15772). It is preferred that you compile from source.

ggml/src/ggml-cpu/arch-fallback.h

@@ -86,8 +86,6 @@
 #define ggml_gemm_mxfp4_8x8_q8_0_generic ggml_gemm_mxfp4_8x8_q8_0
 #define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K
 #elif defined(__x86_64__) || defined(__i386__) || defined(_M_IX86) || defined(_M_X64)
-// quants.c
-#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
 // repack.cpp
 #define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
 #define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4

ggml/src/ggml-cpu/arch/x86/quants.c

@@ -698,6 +698,70 @@ void ggml_vec_dot_q1_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const voi
 #endif
 }
 
+#if defined(__AVX2__)
+// four consecutive Q2_0 weights per byte: byte b becomes b & 3, (b >> 2) & 3, (b >> 4) & 3, (b >> 6) & 3 (OR, never add: the overlapping copies only touch masked-off bits)
+static inline __m256i q2_0_unpack_32_avx2(const uint8_t * qs) {
+    __m256i v = _mm256_cvtepu8_epi32(_mm_loadl_epi64((const __m128i *) qs));
+    v = _mm256_or_si256(v, _mm256_slli_epi32(v, 12));
+    v = _mm256_or_si256(v, _mm256_slli_epi32(v,  6));
+    return _mm256_and_si256(v, _mm256_set1_epi32(0x03030303));
+}
+
+// 32 weights against one Q8_0 block: sum((v - 1) * y) = sum(v * y) - sum(y), exact for every int8 y
+static inline __m256 q2_0_dot_32_avx2(const uint8_t * qs, const int8_t * qy) {
+    const __m256i v = q2_0_unpack_32_avx2(qs);
+    const __m256i y = _mm256_loadu_si256((const __m256i *) qy);
+    const __m256i p = _mm256_sub_epi16(_mm256_maddubs_epi16(v, y), _mm256_maddubs_epi16(_mm256_set1_epi8(1), y));
+    return _mm256_cvtepi32_ps(_mm256_madd_epi16(p, _mm256_set1_epi16(1)));
+}
+
+static inline __m256 q2_0_acc_block_avx2(const block_q2_0 * x, const block_q8_0 * y, __m256 acc) {
+    const float  dx = GGML_CPU_FP16_TO_FP32(x->d);
+    const __m256 d0 = _mm256_set1_ps(dx * GGML_CPU_FP16_TO_FP32(y[0].d));
+    const __m256 d1 = _mm256_set1_ps(dx * GGML_CPU_FP16_TO_FP32(y[1].d));
+    const __m256 p0 = q2_0_dot_32_avx2(x->qs + 0, y[0].qs);
+    const __m256 p1 = q2_0_dot_32_avx2(x->qs + 8, y[1].qs);
+#if defined(__FMA__)
+    acc = _mm256_fmadd_ps(d0, p0, acc);
+    return _mm256_fmadd_ps(d1, p1, acc);
+#else
+    acc = _mm256_add_ps(acc, _mm256_mul_ps(d0, p0));
+    return _mm256_add_ps(acc, _mm256_mul_ps(d1, p1));
+#endif
+}
+#endif
+
+void ggml_vec_dot_q2_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+#if defined(__AVX2__)
+    assert(n % QK2_0 == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q2_0 * GGML_RESTRICT x = vx;
+    const block_q8_0 * GGML_RESTRICT y = vy;
+    const int nb = n / QK2_0;
+
+    // two Q2_0 blocks per step, each into its own accumulator (two FMA chains)
+    __m256 acc0 = _mm256_setzero_ps();
+    __m256 acc1 = _mm256_setzero_ps();
+    int i = 0;
+    for (; i + 1 < nb; i += 2) {
+        acc0 = q2_0_acc_block_avx2(x + i + 0, y + 2*i + 0, acc0);
+        acc1 = q2_0_acc_block_avx2(x + i + 1, y + 2*i + 2, acc1);
+    }
+    if (i < nb) {
+        acc0 = q2_0_acc_block_avx2(x + i, y + 2*i, acc0);
+    }
+
+    *s = hsum_float_8(_mm256_add_ps(acc0, acc1));
+#else
+    ggml_vec_dot_q2_0_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
 void ggml_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
     const int qk = QK8_0;
     const int nb = n / qk;

ggml/src/ggml-cuda/binbcast.cu

@@ -580,8 +580,7 @@ void ggml_cuda_op_repeat_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst
             float       * dst_d  = (float       *) dst->data;
             repeat_back_cuda(src0_d, dst_d, ne00, ne01, ne02, ne03, s00, s01, s02, s03, ne0, ne1, ne2, ne3, stream);
         } break;
-        default: {
+        default:
             GGML_ASSERT(false);
-        } break;
     }
 }

ggml/src/ggml-cuda/common.cuh

@@ -881,6 +881,7 @@ static __device__ __forceinline__ uint8_t ggml_cuda_fp32_to_ue4m3(float x) {
     const __nv_fp8_e4m3 xf(x);
     return xf.__x;
 #else
+     GGML_UNUSED(x);
      NO_DEVICE_CODE; // Used only for NVFP4 Scales for Activations, only for Blackwell
 #endif // defined(BLACKWELL_MMA_AVAILABLE)
 }
@@ -1448,10 +1449,10 @@ struct ggml_backend_cuda_context {
     std::string name;
     cudaEvent_t copy_event = nullptr;
 
-    cudaStream_t streams[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS] = { { nullptr } };
-    cublasHandle_t cublas_handles[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS] = {nullptr};
-    void * cublas_workspaces[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS] = {nullptr};
-    size_t cublas_workspace_sizes[GGML_CUDA_MAX_DEVICES] = {0};
+    cudaStream_t streams[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS] = {};
+    cublasHandle_t cublas_handles[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS] = {};
+    void * cublas_workspaces[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS] = {};
+    size_t cublas_workspace_sizes[GGML_CUDA_MAX_DEVICES] = {};
 
     int curr_stream_no = 0;
 

ggml/src/ggml-cuda/concat.cu

@@ -189,7 +189,6 @@ static void concat_cuda(const ggml_tensor * src0, const ggml_tensor * src1, ggml
                 break;
             default:
                 GGML_ABORT("Invalid dim: %d", dim);
-                break;
         }
     }
 }
@@ -236,7 +235,6 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
                 break;
             default:
                 GGML_ABORT("Unsupported type size: %zu", ggml_type_size(src0->type));
-                break;
         }
     }
 }

ggml/src/ggml-cuda/count-equal.cu

@@ -59,6 +59,5 @@ void ggml_cuda_count_equal(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
         } break;
         default:
             GGML_ASSERT(false);
-            break;
     }
 }

ggml/src/ggml-cuda/fattn-tile.cu

@@ -53,8 +53,7 @@ void ggml_cuda_flash_attn_ext_tile(ggml_backend_cuda_context & ctx, ggml_tensor
             GGML_ASSERT(V->ne[0] == 512);
             ggml_cuda_flash_attn_ext_tile_case<576, 512>(ctx, dst);
         } break;
-        default: {
+        default:
             GGML_ABORT("Unsupported head size");
-        } break;
     }
 }

ggml/src/ggml-cuda/fattn.cu

@@ -412,7 +412,6 @@ static void ggml_cuda_flash_attn_ext_mma_f16(ggml_backend_cuda_context & ctx, gg
         } break;
         default:
             GGML_ABORT("fatal error");
-            break;
     }
 }
 

ggml/src/ggml-cuda/gated_delta_net.cu

@@ -233,7 +233,6 @@ static void launch_gated_delta_net(
         }
         default:
             GGML_ABORT("fatal error");
-            break;
     }
 }
 

ggml/src/ggml-cuda/getrows.cu

@@ -406,7 +406,6 @@ static void ggml_cuda_get_rows_switch_src0_type(
             break;
         default:
             GGML_ABORT("%s: unsupported src0 type: %s\n", __func__, ggml_type_name(src0_type));
-            break;
     }
 }
 
@@ -435,7 +434,6 @@ void get_rows_cuda(
             break;
         default:
             GGML_ABORT("%s: unsupported dst type: %s\n", __func__, ggml_type_name(dst_type));
-            break;
     }
 }
 

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

@@ -138,10 +138,54 @@ int ggml_cuda_get_device() {
     return id;
 }
 
+static bool ggml_cuda_env_flag(const char * name) {
+    const char * env = getenv(name);
+    return env != nullptr && std::atoi(env) != 0;
+}
+
+static int ggml_cuda_env_int(const char * name, int default_value) {
+    const char * env = getenv(name);
+    return env != nullptr ? std::atoi(env) : default_value;
+}
+
+static bool ggml_cuda_env_set(const char * name) {
+    return getenv(name) != nullptr;
+}
+
+static bool ggml_cuda_unified_memory_enabled() {
+    static const bool enabled = ggml_cuda_env_flag("GGML_CUDA_ENABLE_UNIFIED_MEMORY");
+    return enabled;
+}
+
+static bool ggml_cuda_pinned_memory_disabled() {
+    static const bool disabled = ggml_cuda_env_flag("GGML_CUDA_NO_PINNED");
+    return disabled;
+}
+
+static bool ggml_cuda_register_host_enabled() {
+    static const bool enabled = ggml_cuda_env_flag("GGML_CUDA_REGISTER_HOST");
+    return enabled;
+}
+
+static bool ggml_cuda_fusion_disabled() {
+    static const bool disabled = ggml_cuda_env_flag("GGML_CUDA_DISABLE_FUSION");
+    return disabled;
+}
+
+static bool ggml_cuda_p2p_enabled() {
+    static const bool enabled = ggml_cuda_env_set("GGML_CUDA_P2P");
+    return enabled;
+}
+
+static int ggml_cuda_op_offload_min_batch() {
+    static const int min_batch = ggml_cuda_env_int("GGML_OP_OFFLOAD_MIN_BATCH", 32);
+    return min_batch;
+}
+
 static cudaError_t ggml_cuda_device_malloc(void ** ptr, size_t size, int device) {
     ggml_cuda_set_device(device);
     cudaError_t err;
-    if (getenv("GGML_CUDA_ENABLE_UNIFIED_MEMORY") != nullptr) {
+    if (ggml_cuda_unified_memory_enabled()) {
         err = cudaMallocManaged(ptr, size);
 #if defined(GGML_USE_HIP)
         if (err == hipSuccess) {
@@ -383,7 +427,7 @@ static ggml_cuda_device_info ggml_cuda_init() {
     // configure logging to stdout
     // CUBLAS_CHECK(cublasLoggerConfigure(1, 1, 0, nullptr));
 
-    if (getenv("GGML_CUDA_P2P") != nullptr) {
+    if (ggml_cuda_p2p_enabled()) {
         for (int id = 0; id < info.physical_device_count; ++id) {
             CUDA_CHECK(cudaSetDevice(id));
             for (int id_other = 0; id_other < info.physical_device_count; ++id_other) {
@@ -598,7 +642,7 @@ struct ggml_cuda_pool_vmm : public ggml_cuda_pool {
             CU_CHECK(cuMemRelease(handle));
 
             // VMM Bug fix for P2P access if GGML_CUDA_P2P is set, or if NCCL build
-            bool use_peer_access = getenv("GGML_CUDA_P2P") != nullptr;
+            bool use_peer_access = ggml_cuda_p2p_enabled();
 #if defined(GGML_USE_NCCL)
             use_peer_access = true;
 #endif // defined(GGML_USE_NCCL)
@@ -1272,7 +1316,7 @@ static void ggml_backend_cuda_host_buffer_free_buffer(ggml_backend_buffer_t buff
 }
 
 static void * ggml_cuda_host_malloc(size_t size) {
-    if (getenv("GGML_CUDA_NO_PINNED") != nullptr) {
+    if (ggml_cuda_pinned_memory_disabled()) {
         return nullptr;
     }
 
@@ -2641,6 +2685,11 @@ static bool ggml_cuda_is_view_or_noop(const ggml_tensor * t) {
 }
 
 #ifdef USE_CUDA_GRAPH
+static bool ggml_cuda_graph_debug_enabled() {
+    static const bool enabled = ggml_cuda_env_flag("GGML_CUDA_GRAPH_DEBUG");
+    return enabled;
+}
+
 static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) {
 
     bool use_cuda_graph = true;
@@ -2660,9 +2709,9 @@ static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) {
                 // the mul_mat_id fallback path synchronizes the stream, so we cannot use CUDA graphs
                 // ref: https://github.com/ggml-org/llama.cpp/pull/18958
                 use_cuda_graph = false;
-#ifndef NDEBUG
-                GGML_LOG_DEBUG("%s: disabling CUDA graphs due to unsupported node type\n", __func__);
-#endif
+                if (ggml_cuda_graph_debug_enabled()) {
+                    GGML_LOG_DEBUG("%s: disabling CUDA graphs due to unsupported node type\n", __func__);
+                }
             }
         }
 
@@ -2686,7 +2735,9 @@ static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx
 
     if (cgraph->uid != 0 &&
         cgraph->uid == graph->uid) {
-        GGML_LOG_DEBUG("CUDA Graph id %zu reused\n", cgraph->uid);
+        if (ggml_cuda_graph_debug_enabled()) {
+            GGML_LOG_DEBUG("CUDA Graph id %zu reused\n", cgraph->uid);
+        }
         GGML_ASSERT((int)graph->node_props.size() == cgraph->n_nodes);
         return false;
     }
@@ -2733,9 +2784,9 @@ static void ggml_cuda_graph_update_executable(ggml_backend_cuda_context * cuda_c
 #endif // CUDART_VERSION >= 12000
 
     if (stat == cudaErrorGraphExecUpdateFailure) {
-#ifndef NDEBUG
-        GGML_LOG_DEBUG("%s: CUDA graph update failed\n", __func__);
-#endif
+        if (ggml_cuda_graph_debug_enabled()) {
+            GGML_LOG_DEBUG("%s: CUDA graph update failed\n", __func__);
+        }
 
         // The pre-existing graph exec cannot be updated due to violated constraints
         // so instead clear error and re-instantiate
@@ -3605,8 +3656,7 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph *                cgraph,
 // try and fuse nodes and return the number of nodes to skip
 static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) {
 
-    static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION"));
-    if (disable_fusion) {
+    if (ggml_cuda_fusion_disabled()) {
         return 0;
     }
 
@@ -4620,7 +4670,7 @@ static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, co
 
     if (graph->graph == nullptr) {
         if (ggml_cuda_info().devices[cuda_ctx->device].cc < GGML_CUDA_CC_VOLTA) {
-            if (!graph->disable_due_to_gpu_arch) {
+            if (!graph->disable_due_to_gpu_arch && ggml_cuda_graph_debug_enabled()) {
                 GGML_LOG_DEBUG("%s: disabling CUDA graphs due to GPU architecture\n", __func__);
             }
             graph->disable_due_to_gpu_arch = true;
@@ -4655,7 +4705,9 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend,
                 // Warmup: need at least 2 calls with no property change on the 2nd call
                 if (!properties_changed) {
                     graph->warmup_complete = true;
-                    GGML_LOG_DEBUG("%s: CUDA graph warmup complete\n", __func__);
+                    if (ggml_cuda_graph_debug_enabled()) {
+                        GGML_LOG_DEBUG("%s: CUDA graph warmup complete\n", __func__);
+                    }
                     use_cuda_graph = true;
                     cuda_graph_update_required = true;
                 }
@@ -4665,7 +4717,9 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend,
                 if (properties_changed) {
                     // Properties changed - reset warmup, execute directly until stable again
                     graph->warmup_complete = false;
-                    GGML_LOG_DEBUG("%s: CUDA graph warmup reset\n", __func__);
+                    if (ggml_cuda_graph_debug_enabled()) {
+                        GGML_LOG_DEBUG("%s: CUDA graph warmup reset\n", __func__);
+                    }
                 } else {
                     use_cuda_graph = true;
                     cuda_graph_update_required = graph->instance == nullptr;
@@ -4718,7 +4772,7 @@ static void ggml_backend_cuda_event_wait(ggml_backend_t backend, ggml_backend_ev
 static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph, ggml_backend_graph_optimize_params * params) {
     ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context;
 
-    static const bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION"));
+    const bool disable_fusion = ggml_cuda_fusion_disabled();
 
     auto add_alloc_deps = [&](size_t start, size_t last_node) {
 
@@ -5138,7 +5192,7 @@ void ggml_backend_cuda_get_device_memory(int device, size_t * free, size_t * tot
 }
 
 bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size) {
-    if (getenv("GGML_CUDA_REGISTER_HOST") == nullptr) {
+    if (!ggml_cuda_register_host_enabled()) {
         return false;
     }
 
@@ -5161,7 +5215,7 @@ bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size) {
 }
 
 void ggml_backend_cuda_unregister_host_buffer(void * buffer) {
-    if (getenv("GGML_CUDA_REGISTER_HOST") == nullptr) {
+    if (!ggml_cuda_register_host_enabled()) {
         return;
     }
 
@@ -5288,8 +5342,7 @@ static void ggml_backend_cuda_device_get_memory(ggml_backend_dev_t dev, size_t *
     CUDA_CHECK(cudaGetDeviceProperties(&prop, ggml_cuda_get_physical_device(ctx->device)));
 
     // Check if UMA is explicitly enabled via environment variable
-    bool uma_env = getenv("GGML_CUDA_ENABLE_UNIFIED_MEMORY") != nullptr;
-    bool is_uma = prop.integrated > 0 || uma_env;
+    bool is_uma = prop.integrated > 0 || ggml_cuda_unified_memory_enabled();
 
     if (is_uma) {
         // For UMA systems (like DGX Spark), use system memory info
@@ -5330,7 +5383,7 @@ static void ggml_backend_cuda_device_get_props(ggml_backend_dev_t dev, ggml_back
     props->device_id   = ctx->pci_bus_id.empty() ? nullptr : ctx->pci_bus_id.c_str();
     ggml_backend_cuda_device_get_memory(dev, &props->memory_free, &props->memory_total);
 
-    bool host_buffer = getenv("GGML_CUDA_NO_PINNED") == nullptr;
+    const bool host_buffer = !ggml_cuda_pinned_memory_disabled();
 #ifdef GGML_CUDA_NO_PEER_COPY
     bool events = false;
 #else
@@ -5408,7 +5461,6 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                 default:
                     return false;
             }
-            break;
         case GGML_OP_GLU:
             switch (ggml_get_glu_op(op)) {
                 case GGML_GLU_OP_REGLU:
@@ -5426,7 +5478,6 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                 default:
                     return false;
             }
-            break;
         case GGML_OP_MUL_MAT:
         case GGML_OP_MUL_MAT_ID:
             {
@@ -5486,7 +5537,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                     default:
                         return false;
                 }
-            } break;
+            }
         case GGML_OP_OUT_PROD:
             return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32;
         case GGML_OP_GET_ROWS:
@@ -5525,11 +5576,9 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                     default:
                         return false;
                 }
-            } break;
+            }
         case GGML_OP_GET_ROWS_BACK:
-            {
-                return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->ne[2] == 1 && op->ne[3] == 1;
-            } break;
+            return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->ne[2] == 1 && op->ne[3] == 1;
         case GGML_OP_SET_ROWS:
             {
                 return (
@@ -5543,14 +5592,14 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                            )
                        ) &&
                        (op->src[1]->type == GGML_TYPE_I64 || op->src[1]->type == GGML_TYPE_I32);
-            } break;
+            }
         case GGML_OP_SET:
             {
                 const ggml_type t = op->type;
                 return (t == GGML_TYPE_F32 || t == GGML_TYPE_I32) &&
                     t == op->src[0]->type &&
                     t == op->src[1]->type;
-            } break;
+            }
         case GGML_OP_CPY:
             {
                 ggml_type src0_type = op->src[0]->type;
@@ -5606,20 +5655,18 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                     return true;
                 }
                 return false;
-            } break;
+            }
         case GGML_OP_DUP:
                 return true;
         case GGML_OP_ARGMAX:
         case GGML_OP_COUNT_EQUAL:
-            {
-                return true;
-            } break;
+            return true;
         case GGML_OP_REPEAT:
             {
                 // the CUDA REPEAT path only implements F32/F16; other types assert at runtime
                 ggml_type src0_type = op->src[0]->type;
                 return src0_type == GGML_TYPE_F32 || src0_type == GGML_TYPE_F16;
-            } break;
+            }
         case GGML_OP_REPEAT_BACK:
                 return op->type == GGML_TYPE_F32 && (op->src[0]->ne[2]*op->src[0]->ne[3]) <= (1 << 15);
         case GGML_OP_CONCAT:
@@ -5656,7 +5703,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                                )
                            )
                        );
-            } break;
+            }
         case GGML_OP_CONV_TRANSPOSE_1D:
             {
                 ggml_type src0_type = op->src[0]->type;
@@ -5665,7 +5712,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                     return true;
                 }
                 return false;
-            } break;
+            }
         case GGML_OP_COL2IM_1D:
             {
                 ggml_type src0_type = op->src[0]->type;
@@ -5673,17 +5720,15 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                     op->type == src0_type &&
                     ggml_is_contiguous(op->src[0]) &&
                     ggml_is_contiguous(op);
-            } break;
+            }
         case GGML_OP_SILU_BACK:
             return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32;
-            break;
         case GGML_OP_NORM:
         case GGML_OP_RMS_NORM:
         case GGML_OP_L2_NORM:
             return ggml_is_contiguous_rows(op->src[0]);
         case GGML_OP_RMS_NORM_BACK:
             return ggml_is_contiguous(op->src[0]);
-            break;
         case GGML_OP_NONE:
         case GGML_OP_RESHAPE:
         case GGML_OP_VIEW:
@@ -6034,7 +6079,7 @@ ggml_backend_reg_t ggml_backend_cuda_reg() {
         std::lock_guard<std::mutex> lock(mutex);
         if (!initialized) {
             ggml_backend_cuda_reg_context * ctx = new ggml_backend_cuda_reg_context;
-            const int min_batch_size = getenv("GGML_OP_OFFLOAD_MIN_BATCH") ? atoi(getenv("GGML_OP_OFFLOAD_MIN_BATCH")) : 32;
+            const int min_batch_size = ggml_cuda_op_offload_min_batch();
 
             const ggml_cuda_device_info & info = ggml_cuda_info();
             const bool virtual_devices = info.device_count > info.physical_device_count;

ggml/src/ggml-cuda/mmf.cuh

@@ -741,9 +741,8 @@ void mul_mat_f_cuda(
                 sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, block_nums, block_dims, nbytes_shared_total, stream,
                 ids_data);
         } break;
-        default: {
+        default:
             GGML_ABORT("fatal error");
-        } break;
     }
 
     GGML_UNUSED_VARS(nchannels_y);
@@ -845,9 +844,8 @@ static void mul_mat_f_switch_cols_per_block(
                 stride_col_id, stride_row_id, nchannels_x, nchannels_y,  nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst,
                 nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, stream, ids_data);
         } break;
-        default: {
+        default:
             GGML_ABORT("fatal error");
-        } break;
     }
 }
 

ggml/src/ggml-cuda/mmq-load-tiles.cuh

@@ -9,9 +9,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
     int   * x_qs = (int   *)  x_tile;
     float * x_df = (float *) (x_qs + 2*MMQ_TILE_NE_K);
 #else
@@ -100,9 +100,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
     int   * x_qs = (int   *)  x_tile;
     float * x_df = (float *) (x_qs + 2*MMQ_TILE_NE_K);
 #else
@@ -189,9 +189,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
     int   * x_qs = (int   *)  x_tile;
     float * x_df = (float *) (x_qs + 2*MMQ_TILE_NE_K);
 #else
@@ -252,9 +252,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
     int   * x_qs = (int   *)  x_tile;
     half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K);
 #else
@@ -315,9 +315,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
     int   * x_qs = (int   *)  x_tile;
     float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
 #else
@@ -395,9 +395,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
     int   * x_qs = (int   *)  x_tile;
     half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K);
 #else
@@ -473,9 +473,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
     int   * x_qs = (int   *)  x_tile;
     float * x_df = (float *) (x_tile + 2*MMQ_TILE_NE_K);
 #else
@@ -539,9 +539,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
     int   * x_qs = (int   *)  x_tile;
     half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K);
 #else
@@ -600,9 +600,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
     int   * x_qs = (int   *)  x_tile;
     float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
 #else
@@ -713,9 +713,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
     int   * x_qs = (int   *)  x_tile;
     half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K);
 #else
@@ -780,7 +780,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
             const half2 dm = bxi->dm * make_half2(1.0f, -1.0f);
 
     #pragma unroll
-            for (int l = 0; l < sizeof(int); ++l) {
+            for (int l = 0; l < int(sizeof(int)); ++l) {
                 x_dm[i*sram_stride + sizeof(int)*ksc + l] = dm*make_half2(sc8[l], m8[l]);
             }
         }
@@ -824,9 +824,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
     int   * x_qs = (int   *)  x_tile;
     half2 * x_dm = (half2 *) (x_qs + MMQ_TILE_NE_K*2);
 #else
@@ -948,9 +948,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
     int   * x_qs = (int   *)  x_tile;
     float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
     int   * x_sc = (int   *) (x_df + MMQ_TILE_NE_K/QI6_K);
@@ -1038,9 +1038,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
     int   * x_qs = (int   *)  x_tile;
     half2 * x_ds = (half2 *) (x_qs + MMQ_TILE_NE_K*2);
 #else
@@ -1100,9 +1100,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
     int   * x_qs = (int   *)  x_tile;
     float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
 #else
@@ -1164,9 +1164,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
     int   * x_qs = (int   *)  x_tile;
     float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
 #else
@@ -1229,9 +1229,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
     int   * x_qs = (int   *)  x_tile;
     float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
 #else
@@ -1297,9 +1297,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
     int   * x_qs = (int   *)  x_tile;
     float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
 #else
@@ -1361,9 +1361,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
     int   * x_qs = (int   *)  x_tile;
     float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
 #else
@@ -1430,9 +1430,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
     int   * x_qs = (int   *)  x_tile;
     float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
 #else
@@ -1497,9 +1497,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
     int   * x_qs = (int   *)  x_tile;
     float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
 #else
@@ -1566,9 +1566,9 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, GGML_PREC_Q8) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, GGML_PREC_Q8);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, GGML_PREC_Q8);
     int   * x_qs = (int   *)  x_tile;
     float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
 #else
@@ -1675,9 +1675,9 @@ template <ggml_type type, int J, bool fallback, ggml_prec prec_src1> static __de
     constexpr int warp_size   = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps      = ggml_cuda_mmq_get_nthreads(type, J, fallback, prec_src1) / warp_size;
     constexpr int I           = ggml_cuda_mmq_get_I(type, J, fallback, prec_src1);
-    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, prec_src1);
 
 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
+    constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback, prec_src1);
     int   * x_qs = (int   *) x_tile;
     float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2);
 #else

ggml/src/ggml-cuda/mmq.cu

@@ -87,7 +87,6 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con
             break;
         default:
             GGML_ABORT("fatal error");
-            break;
     }
 }
 

ggml/src/ggml-cuda/mmq.cuh

@@ -96,7 +96,6 @@ static mmq_q8_1_ds_layout mmq_get_q8_1_ds_layout(const ggml_type type_x) {
             return MMQ_Q8_1_DS_LAYOUT_D4;
         default:
             GGML_ABORT("fatal error");
-            break;
     }
 }
 
@@ -1529,7 +1528,6 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args,
         default:
             fprintf(stderr, "J_best=%d\n", args.J_best);
             GGML_ABORT("fatal error");
-            break;
     }
 }
 

ggml/src/ggml-cuda/mmvf.cu

@@ -504,9 +504,8 @@ void launch_mul_mat_vec_f_cuda(
                  channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst,
                  sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, block_dims, block_nums, nbytes_shared, ids_stride, stream);
         } break;
-        default: {
+        default:
             GGML_ABORT("fatal error");
-        } break;
     }
 }
 
@@ -603,7 +602,6 @@ static void mul_mat_vec_f_cuda_switch_ncols_dst(
             break;
         default:
             GGML_ABORT("fatal error");
-            break;
     }
 }
 

ggml/src/ggml-cuda/mmvq.cu

@@ -1275,7 +1275,6 @@ static void mul_mat_vec_q_switch_ncols_dst(
         } break;
         default:
             GGML_ABORT("fatal error");
-            break;
     }
 }
 static void mul_mat_vec_q_switch_type(
@@ -1427,7 +1426,6 @@ static void mul_mat_vec_q_switch_type(
             break;
         default:
             GGML_ABORT("fatal error");
-            break;
     }
 }
 

ggml/src/ggml-cuda/quantize.cu

@@ -597,7 +597,6 @@ void quantize_mmq_q8_1_cuda(
             break;
         default:
             GGML_ABORT("fatal error");
-            break;
     }
 }
 
@@ -627,7 +626,6 @@ void quantize_scatter_mmq_q8_1_cuda(
             break;
         default:
             GGML_ABORT("fatal error");
-            break;
     }
 }
 

ggml/src/ggml-cuda/topk-moe.cu

@@ -348,7 +348,6 @@ static void launch_topk_moe_cuda(ggml_backend_cuda_context & ctx,
             break;
         default:
             GGML_ASSERT(false && "fatal error");
-            break;
     }
 }
 

ggml/src/ggml-hexagon/ggml-hexagon.cpp

@@ -344,7 +344,8 @@ static bool ggml_hexagon_precompute_binary_params(
     const struct ggml_tensor * src0,
     const struct ggml_tensor * src1,
     const struct ggml_tensor * dst,
-    struct htp_binary_kernel_params * kparams
+    struct htp_binary_kernel_params * kparams,
+    const struct ggml_tensor * src2 = nullptr
 );
 
 static void ggml_hexagon_precompute_get_rows_params(
@@ -5864,7 +5865,8 @@ static bool ggml_hexagon_precompute_binary_params(
     const struct ggml_tensor * src0,
     const struct ggml_tensor * src1,
     const struct ggml_tensor * dst,
-    struct htp_binary_kernel_params * kparams
+    struct htp_binary_kernel_params * kparams,
+    const struct ggml_tensor * src2
 ) {
     memset(kparams, 0, sizeof(*kparams));
 
@@ -5877,14 +5879,17 @@ static bool ggml_hexagon_precompute_binary_params(
     const size_t src1_row_size_aligned = hex_round_up(src1_row_size, 128);
     const size_t dst_row_size_aligned  = hex_round_up(dst_row_size,  128);
 
+    if (!htp_binary_strides_safe(src0, src1, dst, elem_size)) {
+        return false;
+    }
+
     const bool is_add_id = op == HTP_OP_ADD_ID;
     const bool is_scalar = !is_add_id && src1->ne[0] == 1;
-    const bool is_transposed = src0->nb[1] < src0_row_size || src1->nb[1] < src1_row_size || dst->nb[1] < dst_row_size;
-    const bool is_row_bcast = !is_add_id && !is_scalar && !is_transposed &&
+    const bool is_row_bcast = !is_add_id && !is_scalar &&
         src1->ne[0] == src0->ne[0] &&
         (src0->ne[1] > 1 || src0->ne[2] > 1 || src0->ne[3] > 1) &&
         src1->ne[1] == 1 && src1->ne[2] == 1 && src1->ne[3] == 1;
-    const bool is_same_shape = !is_add_id && !is_scalar && !is_transposed &&
+    const bool is_same_shape = !is_add_id && !is_scalar &&
         src1->ne[0] == src0->ne[0] &&
         src1->ne[1] == src0->ne[1] &&
         (src1->ne[2] == src0->ne[2] || src1->ne[2] == 1) &&
@@ -5925,28 +5930,42 @@ static bool ggml_hexagon_precompute_binary_params(
 
     enum htp_binary_kernel_type kernel_type;
     size_t src1_size = 0;
+    size_t src2_size = 0;
 
     if (is_add_id) {
+        GGML_ASSERT(src2 != nullptr);
         kernel_type = HTP_BINARY_KERNEL_ADD_ID;
         src1_size = hex_round_up(src1->ne[1] * src1_row_size_aligned, 128);
+        src2_size = hex_round_up(ggml_nelements(src2) * sizeof(int32_t), 128);
+        if (src1_size + src2_size > sess->vtcm_size) {
+            return false;
+        }
     } else if (is_row_bcast) {
         kernel_type = HTP_BINARY_KERNEL_ROW_BCAST;
         src1_size = src1_row_size_aligned;
     } else if (is_scalar) {
-        const bool is_scalar_static = (src1->ne[2] == 1 && src1->ne[3] == 1) &&
-            (src1->ne[1] == 1 || src1->nb[1] == elem_size);
-        if (is_scalar_static) {
-            kernel_type = HTP_BINARY_KERNEL_SCALAR_DMA;
-            src1_size = hex_round_up(src1->ne[1] * elem_size, 128);
-        } else {
-            kernel_type = HTP_BINARY_KERNEL_SCALAR;
+        src1_size = hex_round_up(ggml_nelements(src1) * elem_size, 128);
+        if (src1_size > sess->vtcm_size) {
+            return false;
         }
+        kernel_type = HTP_BINARY_KERNEL_SCALAR;
+        kparams->src1_strided = !ggml_is_contiguous(src1) ? 1 : 0;
     } else if (is_same_shape) {
         kernel_type = HTP_BINARY_KERNEL_SAME_SHAPE;
     } else if (is_complex) {
+        src1_size = hex_round_up(ggml_nelements(src1) * elem_size + 128, 128);
+        if (src1_size > sess->vtcm_size) {
+            return false;
+        }
         kernel_type = HTP_BINARY_KERNEL_COMPLEX;
+        kparams->src1_strided = !ggml_is_contiguous(src1) ? 1 : 0;
     } else {
+        src1_size = hex_round_up(ggml_nelements(src1) * elem_size + 128, 128);
+        if (src1_size > sess->vtcm_size) {
+            return false;
+        }
         kernel_type = HTP_BINARY_KERNEL_REPEAT;
+        kparams->src1_strided = !ggml_is_contiguous(src1) ? 1 : 0;
     }
 
     kparams->kernel_type = kernel_type;
@@ -5955,6 +5974,39 @@ static bool ggml_hexagon_precompute_binary_params(
     kparams->src1_row_size_aligned = src1_row_size_aligned;
     kparams->dst_row_size_aligned  = dst_row_size_aligned;
     kparams->src1_size             = src1_size;
+    kparams->src2_size             = src2_size;
+
+    const bool src0_contig_dim1 = (src0->nb[2] == src0->ne[1] * src0->nb[1]);
+    const bool src1_contig_dim1 = (src1->nb[2] == src1->ne[1] * src1->nb[1]);
+    const bool dst_contig_dim1  = (dst->nb[2]  == src0->ne[1] * dst->nb[1]);
+
+    const bool src0_contig_dim2 = (src0->nb[3] == src0->ne[2] * src0->nb[2]);
+    const bool src1_contig_dim2 = (src1->nb[3] == src1->ne[2] * src1->nb[2]);
+    const bool dst_contig_dim2  = (dst->nb[3]  == src0->ne[2] * dst->nb[2]);
+
+    bool split_at_ne01 = false;
+    if (src0->ne[2] > 1) {
+        if (kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE) {
+            split_at_ne01 = (src1->ne[2] != src0->ne[2]) || !src0_contig_dim1 || !src1_contig_dim1 || !dst_contig_dim1;
+        } else if (kernel_type == HTP_BINARY_KERNEL_ROW_BCAST) {
+            split_at_ne01 = !src0_contig_dim1 || !dst_contig_dim1;
+        } else {
+            split_at_ne01 = true;
+        }
+    }
+    kparams->split_at_ne01 = split_at_ne01 ? 1 : 0;
+
+    bool split_at_ne02 = false;
+    if (src0->ne[3] > 1) {
+        if (kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE) {
+            split_at_ne02 = (src1->ne[3] != src0->ne[3]) || !src0_contig_dim2 || !src1_contig_dim2 || !dst_contig_dim2;
+        } else if (kernel_type == HTP_BINARY_KERNEL_ROW_BCAST) {
+            split_at_ne02 = !src0_contig_dim2 || !dst_contig_dim2;
+        } else {
+            split_at_ne02 = true;
+        }
+    }
+    kparams->split_at_ne02 = split_at_ne02 ? 1 : 0;
 
     struct htp_binary_vtcm_layout L;
     htp_binary_vtcm_layout_build(&L, kparams, sess->vtcm_size);
@@ -7139,9 +7191,12 @@ static bool ggml_hexagon_supported_add_id(const struct ggml_hexagon_session * se
     if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(src1) || !ggml_is_contiguous(dst)) {
         return false;
     }
+    if (src2->nb[1] < src2->ne[0] * sizeof(int32_t)) {
+        return false;
+    }
 
     struct htp_binary_kernel_params kparams;
-    return ggml_hexagon_precompute_binary_params(sess, HTP_OP_ADD_ID, src0, src1, dst, &kparams);
+    return ggml_hexagon_precompute_binary_params(sess, HTP_OP_ADD_ID, src0, src1, dst, &kparams, src2);
 }
 
 static bool ggml_hexagon_supported_unary(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
@@ -8504,9 +8559,10 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg
                 );
             } else if (node.opcode == HTP_OP_MUL || node.opcode == HTP_OP_ADD || node.opcode == HTP_OP_ADD_ID || node.opcode == HTP_OP_SUB || node.opcode == HTP_OP_DIV) {
                 const ggml_tensor * src1 = node.node->src[1];
+                const ggml_tensor * src2 = (node.opcode == HTP_OP_ADD_ID) ? node.node->src[2] : nullptr;
                 GGML_ASSERT(ggml_hexagon_precompute_binary_params(sess,
                     node.opcode, node.node->src[0], src1, node.node,
-                    (struct htp_binary_kernel_params *) node.kernel_params));
+                    (struct htp_binary_kernel_params *) node.kernel_params, src2));
             } else if (node.opcode == HTP_OP_FLASH_ATTN_EXT) {
                 ggml_hexagon_precompute_flash_attn_params(sess,
                     node.node,
@@ -9391,9 +9447,11 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
                 case GGML_GLU_OP_SWIGLU_OAI:
                 case GGML_GLU_OP_SWIGLU_CLAMP:
                 case GGML_GLU_OP_GEGLU:
+                    supp = ggml_hexagon_supported_activations(sess, op);
+                    break;
                 case GGML_GLU_OP_GEGLU_QUICK:
                 case GGML_GLU_OP_GEGLU_ERF:
-                    supp = ggml_hexagon_supported_activations(sess, op);
+                    supp = (op->src[0]->type == GGML_TYPE_F32) && ggml_hexagon_supported_activations(sess, op);
                     break;
                 default:
                     supp = false;

ggml/src/ggml-hexagon/htp/binary-ops.c

@@ -103,13 +103,13 @@ static inline uint32_t calc_block_size(struct htp_binary_context * bctx, uint32_
 
 // Out-of-line compute micro-kernels
 
-typedef void (*compute_scalar_dma_t)(
+typedef void (*compute_scalar_t)(
     uint8_t * dst, const uint8_t * src0, const void * s1_table,
     uint32_t cur_i11, uint32_t ne11, uint32_t n_rows,
     size_t dst_stride, size_t src0_stride, uint32_t ne00);
 
-#define DEFINE_COMPUTE_SCALAR_DMA(NAME, TYPE, HVX_STMT)         \
-static void compute_scalar_dma_##NAME(                          \
+#define DEFINE_COMPUTE_SCALAR(NAME, TYPE, HVX_STMT)             \
+static void compute_scalar_##NAME(                              \
     uint8_t * dst, const uint8_t * src0, const void * s1_table, \
     uint32_t cur_i11, uint32_t ne11, uint32_t n_rows,           \
     size_t dst_stride, size_t src0_stride, uint32_t ne00) {     \
@@ -125,31 +125,6 @@ static void compute_scalar_dma_##NAME(                          \
     }                                                           \
 }
 
-DEFINE_COMPUTE_SCALAR_DMA(add_f32, float,    hvx_add_scalar_f32_aa(r_dst, r_src0, val, ne00))
-DEFINE_COMPUTE_SCALAR_DMA(add_f16, _Float16, hvx_add_scalar_f16_aa(r_dst, r_src0, val, ne00))
-DEFINE_COMPUTE_SCALAR_DMA(sub_f32, float,    hvx_sub_scalar_f32_aa(r_dst, r_src0, val, ne00))
-DEFINE_COMPUTE_SCALAR_DMA(sub_f16, _Float16, hvx_sub_scalar_f16_aa(r_dst, r_src0, val, ne00))
-DEFINE_COMPUTE_SCALAR_DMA(mul_f32, float,    hvx_mul_scalar_f32_aa(r_dst, r_src0, val, ne00))
-DEFINE_COMPUTE_SCALAR_DMA(mul_f16, _Float16, hvx_mul_scalar_f16_aa(r_dst, r_src0, val, ne00))
-DEFINE_COMPUTE_SCALAR_DMA(div_f32, float,    hvx_mul_scalar_f32_aa(r_dst, r_src0, 1.0f / (val), ne00))
-DEFINE_COMPUTE_SCALAR_DMA(div_f16, _Float16, hvx_div_scalar_f16_aa(r_dst, r_src0, val, ne00))
-
-typedef void (*compute_scalar_t)(
-    uint8_t * dst, const uint8_t * src0, const uint8_t * src1_ptr, uint32_t s1_stride,
-    uint32_t n_rows, size_t dst_stride, size_t src0_stride, uint32_t ne00);
-
-#define DEFINE_COMPUTE_SCALAR(NAME, TYPE, HVX_STMT)                                          \
-static void compute_scalar_##NAME(                                                           \
-    uint8_t * dst, const uint8_t * src0, const uint8_t * src1_ptr, uint32_t s1_stride,       \
-    uint32_t n_rows, size_t dst_stride, size_t src0_stride, uint32_t ne00) {                 \
-    for (uint32_t r = 0; r < n_rows; r++) {                                                  \
-        uint8_t * r_dst = dst + r * dst_stride;                                              \
-        const uint8_t * r_src0 = src0 + r * src0_stride;                                     \
-        TYPE val = *(const TYPE *)(src1_ptr + r * s1_stride);                                \
-        HVX_STMT;                                                                            \
-    }                                                                                        \
-}
-
 DEFINE_COMPUTE_SCALAR(add_f32, float,    hvx_add_scalar_f32_aa(r_dst, r_src0, val, ne00))
 DEFINE_COMPUTE_SCALAR(add_f16, _Float16, hvx_add_scalar_f16_aa(r_dst, r_src0, val, ne00))
 DEFINE_COMPUTE_SCALAR(sub_f32, float,    hvx_sub_scalar_f32_aa(r_dst, r_src0, val, ne00))
@@ -205,17 +180,17 @@ DEFINE_COMPUTE_ROW_BCAST(div_f16, hvx_div_f16_aaa)
 
 typedef void (*compute_complex_t)(
     uint8_t * dst, const uint8_t * src0, const uint8_t * src1_plane,
-    uint32_t i01, uint32_t ne11, const struct fastdiv_values * div11, uint32_t nb11,
+    uint32_t i01, uint32_t ne11, const struct fastdiv_values * div11, uint32_t src1_stride,
     uint32_t n_rows, size_t dst_stride, size_t src0_stride, uint32_t ne00);
 
 #define DEFINE_COMPUTE_COMPLEX(NAME, HVX_FN)                                                 \
 static void compute_complex_##NAME(                                                          \
     uint8_t * dst, const uint8_t * src0, const uint8_t * src1_plane,                         \
-    uint32_t i01, uint32_t ne11, const struct fastdiv_values * div11, uint32_t nb11,         \
+    uint32_t i01, uint32_t ne11, const struct fastdiv_values * div11, uint32_t src1_stride,  \
     uint32_t n_rows, size_t dst_stride, size_t src0_stride, uint32_t ne00) {                 \
     for (uint32_t r = 0; r < n_rows; r++) {                                                  \
         uint32_t i11 = fastmodulo(i01 + r, ne11, div11);                                     \
-        const uint8_t * r_src1 = src1_plane + i11 * nb11;                                    \
+        const uint8_t * r_src1 = src1_plane + i11 * src1_stride;                             \
         HVX_FN(dst + r * dst_stride, src0 + r * src0_stride, r_src1, ne00);                  \
     }                                                                                        \
 }
@@ -231,17 +206,17 @@ DEFINE_COMPUTE_COMPLEX(div_f16, hvx_div_f16_aau)
 
 typedef void (*compute_repeat_t)(
     uint8_t * dst, const uint8_t * src0, const uint8_t * src1_plane,
-    uint32_t i01, uint32_t ne11, const struct fastdiv_values * div11, uint32_t nb11,
+    uint32_t i01, uint32_t ne11, const struct fastdiv_values * div11, uint32_t src1_stride,
     uint32_t n_rows, size_t dst_stride, size_t src0_stride, uint32_t ne00, uint32_t ne10);
 
 #define DEFINE_COMPUTE_REPEAT(NAME, TYPE, HVX_FN)                                            \
 static void compute_repeat_##NAME(                                                           \
     uint8_t * dst, const uint8_t * src0, const uint8_t * src1_plane,                         \
-    uint32_t i01, uint32_t ne11, const struct fastdiv_values * div11, uint32_t nb11,         \
+    uint32_t i01, uint32_t ne11, const struct fastdiv_values * div11, uint32_t src1_stride,  \
     uint32_t n_rows, size_t dst_stride, size_t src0_stride, uint32_t ne00, uint32_t ne10) {  \
     for (uint32_t r = 0; r < n_rows; r++) {                                                  \
         uint32_t i11 = fastmodulo(i01 + r, ne11, div11);                                     \
-        const uint8_t * r_src1_row = src1_plane + i11 * nb11;                                \
+        const uint8_t * r_src1_row = src1_plane + i11 * src1_stride;                         \
         uint8_t * r_dst = dst + r * dst_stride;                                              \
         const uint8_t * r_src0 = src0 + r * src0_stride;                                     \
         for (uint32_t c = 0; c < ne00; c += ne10) {                                          \
@@ -261,17 +236,17 @@ DEFINE_COMPUTE_REPEAT(div_f32, float,    hvx_div_f32_uuu)
 DEFINE_COMPUTE_REPEAT(div_f16, _Float16, hvx_div_f16_uuu)
 
 typedef void (*compute_add_id_t)(
-    uint8_t * dst, const uint8_t * src0, const uint8_t * src1_data, const char * src2_data,
-    uint32_t i01, uint32_t i02, uint32_t nb20, uint32_t nb21, uint32_t src1_stride,
+    uint8_t * dst, const uint8_t * src0, const uint8_t * src1_data, const int32_t * src2_ids,
+    uint32_t i01, uint32_t i02, uint32_t ne01, uint32_t src1_stride,
     uint32_t n_rows, size_t dst_stride, size_t src0_stride, uint32_t ne00);
 
 static void compute_add_id_f32(
-    uint8_t * dst, const uint8_t * src0, const uint8_t * src1_data, const char * src2_data,
-    uint32_t i01, uint32_t i02, uint32_t nb20, uint32_t nb21, uint32_t src1_stride,
+    uint8_t * dst, const uint8_t * src0, const uint8_t * src1_data, const int32_t * src2_ids,
+    uint32_t i01, uint32_t i02, uint32_t ne01, uint32_t src1_stride,
     uint32_t n_rows, size_t dst_stride, size_t src0_stride, uint32_t ne00) {
     for (uint32_t r = 0; r < n_rows; r++) {
         uint32_t r_i01 = i01 + r;
-        const int32_t idx = *(const int32_t *)(src2_data + r_i01 * nb20 + i02 * nb21);
+        const int32_t idx = src2_ids[i02 * ne01 + r_i01];
         if (idx < 0) {
             memcpy(dst + r * dst_stride, src0 + r * src0_stride, ne00 * sizeof(float));
             continue;
@@ -283,18 +258,22 @@ static void compute_add_id_f32(
     }
 }
 
-// 1a. Scalar src1 in VTCM via DMA (ne10 == 1, ne12 == 1, ne13 == 1)
-static void binary_thread_scalar_dma(unsigned int nth, unsigned int ith, void * data) {
+// 1. Scalar src1 in VTCM via DMA (ne10 == 1)
+static void binary_thread_scalar(unsigned int nth, unsigned int ith, void * data) {
     struct htp_binary_context * bctx = (struct htp_binary_context *) data;
     struct htp_ops_context * octx = bctx->octx;
     htp_binary_preamble;
 
+    const uint32_t elem_size = (src0->type == HTP_TYPE_F32) ? sizeof(float) : sizeof(_Float16);
+    const size_t plane_stride = (size_t) ne12 * ne11 * elem_size;
+    const size_t row_stride   = (size_t) ne11 * elem_size;
+
     const uint32_t row_size_bytes = bctx->row_size_bytes;
     const uint32_t start_row = bctx->row_start + bctx->nrows_per_thread * ith;
     const uint32_t end_row   = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
     if (start_row >= end_row) return;
 
-    FARF(HIGH, "binary-scalar-dma: %d/%d (%u:%u) row-size %u (%u)",
+    FARF(HIGH, "binary-scalar: %d/%d (%u:%u) row-size %u (%u)",
          ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned);
 
     const struct htp_binary_vtcm_layout * layout = &bctx->vtcm_layout;
@@ -329,94 +308,6 @@ static void binary_thread_scalar_dma(unsigned int nth, unsigned int ith, void *
         spad_idx ^= 1;
     }
 
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
-    compute_scalar_dma_t compute = (compute_scalar_dma_t) bctx->compute;
-
-    for (uint32_t ir = start_row; ir < end_row; ) {
-        uint32_t current_block_size = calc_block_size(bctx, ir, end_row, ne01, ne02);
-
-        uint8_t * d_spad  = (uint8_t *) dma_queue_pop(dma_q).src;
-        uint8_t * s0_spad = (uint8_t *) dma_queue_pop(dma_q).dst;
-
-        uint32_t i03, i02, i01, rem;
-        i03 = fastdiv(ir, &bctx->src0_dim12_div);
-        rem = ir - i03 * (ne02 * ne01);
-        i02 = fastdiv(rem, &bctx->src0_dim1_div);
-        i01 = rem - i02 * ne01;
-
-        uint32_t cur_i11 = fastmodulo(i01, ne11, &bctx->src1_dim1_div);
-
-        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
-        compute(d_spad, s0_spad, s1_table, cur_i11, ne11, current_block_size,
-                bctx->dst_row_size_aligned, bctx->src0_row_size_aligned, ne00);
-        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
-
-        dma_addr_t dst_curr = dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1;
-        dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes,
-                                                                current_block_size);
-
-        if (ir_prefetch < end_row) {
-            uint32_t next_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02);
-            uint32_t p03, p02, p01, prem;
-            p03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div);
-            prem = ir_prefetch - p03 * (ne02 * ne01);
-            p02 = fastdiv(prem, &bctx->src0_dim1_div);
-            p01 = prem - p02 * ne01;
-            dma_addr_t s0_next = src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01;
-            dma_queue_push(dma_q, dma_make_data(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes,
-                                                               next_block_size);
-            ir_prefetch += next_block_size;
-        }
-        ir += current_block_size;
-    }
-
-    dma_queue_flush(dma_q);
-}
-
-// 1b. Scalar src1 dynamic / pointer (ne10 == 1)
-static void binary_thread_scalar(unsigned int nth, unsigned int ith, void * data) {
-    struct htp_binary_context * bctx = (struct htp_binary_context *) data;
-    struct htp_ops_context * octx = bctx->octx;
-    htp_binary_preamble;
-
-    const uint32_t row_size_bytes = bctx->row_size_bytes;
-    const uint32_t start_row = bctx->row_start + bctx->nrows_per_thread * ith;
-    const uint32_t end_row   = MIN(start_row + bctx->nrows_per_thread, bctx->row_start + bctx->total_rows);
-    if (start_row >= end_row) return;
-
-    FARF(HIGH, "binary-scalar: %d/%d (%u:%u) row-size %u (%u)",
-         ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned);
-
-    const struct htp_binary_vtcm_layout * layout = &bctx->vtcm_layout;
-    uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_thread);
-    uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_thread);
-    size_t src0_spad_half    = layout->src0_spad_half_size;
-    size_t dst_spad_half     = layout->dst_spad_half_size;
-
-    dma_queue * dma_q = octx->ctx->dma[ith];
-    uint32_t ir_prefetch = start_row;
-    int spad_idx = 0;
-
-    for (int k = 0; k < 2 && ir_prefetch < end_row; k++) {
-        uint32_t current_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02);
-        uint32_t i03, i02, i01, rem;
-        i03 = fastdiv(ir_prefetch, &bctx->src0_dim12_div);
-        rem = ir_prefetch - i03 * (ne02 * ne01);
-        i02 = fastdiv(rem, &bctx->src0_dim1_div);
-        i01 = rem - i02 * ne01;
-
-        dma_addr_t src0_curr = src0->data + i03 * nb03 + i02 * nb02 + i01 * nb01;
-        dma_addr_t dst_curr  = dst->data  + i03 * nb3  + i02 * nb2  + i01 * nb1;
-
-        uint8_t * s0_spad = src0_spad_base + spad_idx * src0_spad_half;
-        uint8_t * d_spad  = dst_spad_base  + spad_idx * dst_spad_half;
-
-        dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, 0);
-        dma_queue_push(dma_q, dma_make_data(s0_spad, src0_curr), bctx->src0_row_size_aligned, nb01, row_size_bytes, current_block_size);
-        ir_prefetch += current_block_size;
-        spad_idx ^= 1;
-    }
-
     struct htp_thread_trace * tr = &octx->ctx->trace[ith];
     compute_scalar_t compute = (compute_scalar_t) bctx->compute;
 
@@ -435,17 +326,17 @@ static void binary_thread_scalar(unsigned int nth, unsigned int ith, void * data
         uint32_t i13 = fastmodulo(i03, ne13, &bctx->src1_dim3_div);
         uint32_t i12 = fastmodulo(i02, ne12, &bctx->src1_dim2_div);
         uint32_t i11 = fastmodulo(i01, ne11, &bctx->src1_dim1_div);
-
-        const uint8_t * src1_ptr = (const uint8_t *)(uintptr_t) src1->data + i13 * nb13 + i12 * nb12 + i11 * nb11;
-        uint32_t s1_stride = (ne11 == 1) ? 0 : nb11;
+        const uint8_t * scalar_table = (const uint8_t *) s1_table +
+            i13 * plane_stride + i12 * row_stride;
 
         htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
-        compute(d_spad, s0_spad, src1_ptr, s1_stride, current_block_size,
+        compute(d_spad, s0_spad, scalar_table, i11, ne11, current_block_size,
                 bctx->dst_row_size_aligned, bctx->src0_row_size_aligned, ne00);
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
 
         dma_addr_t dst_curr = dst->data + i03 * nb3 + i02 * nb2 + i01 * nb1;
-        dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes, current_block_size);
+        dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), nb1, bctx->dst_row_size_aligned, row_size_bytes,
+                                                                current_block_size);
 
         if (ir_prefetch < end_row) {
             uint32_t next_block_size = calc_block_size(bctx, ir_prefetch, end_row, ne01, ne02);
@@ -455,7 +346,8 @@ static void binary_thread_scalar(unsigned int nth, unsigned int ith, void * data
             p02 = fastdiv(prem, &bctx->src0_dim1_div);
             p01 = prem - p02 * ne01;
             dma_addr_t s0_next = src0->data + p03 * nb03 + p02 * nb02 + p01 * nb01;
-            dma_queue_push(dma_q, dma_make_data(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes, next_block_size);
+            dma_queue_push(dma_q, dma_make_data(s0_spad, s0_next), bctx->src0_row_size_aligned, nb01, row_size_bytes,
+                                                               next_block_size);
             ir_prefetch += next_block_size;
         }
         ir += current_block_size;
@@ -668,6 +560,12 @@ static void binary_thread_vector_complex(unsigned int nth, unsigned int ith, voi
     uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_thread);
     size_t src0_spad_half    = layout->src0_spad_half_size;
     size_t dst_spad_half     = layout->dst_spad_half_size;
+    const uint8_t * vtcm_src1 = VTCM_LAYOUT_PTR(const uint8_t, bctx->vtcm_base, layout->off_src1);
+
+    const uint32_t elem_size = (src0->type == HTP_TYPE_F32) ? sizeof(float) : sizeof(_Float16);
+    const size_t vtcm_row_stride   = (size_t) ne10 * elem_size;
+    const size_t vtcm_plane_stride = (size_t) ne11 * vtcm_row_stride;
+    const size_t vtcm_vol_stride   = (size_t) ne12 * vtcm_plane_stride;
 
     dma_queue * dma_q = octx->ctx->dma[ith];
     uint32_t ir_prefetch = start_row;
@@ -707,10 +605,10 @@ static void binary_thread_vector_complex(unsigned int nth, unsigned int ith, voi
 
         uint32_t i13 = fastmodulo(i03, ne13, &bctx->src1_dim3_div);
         uint32_t i12 = fastmodulo(i02, ne12, &bctx->src1_dim2_div);
-        const uint8_t * src1_plane = (const uint8_t *)(uintptr_t) src1->data + i13 * nb13 + i12 * nb12;
+        const uint8_t * src1_plane = vtcm_src1 + i13 * vtcm_vol_stride + i12 * vtcm_plane_stride;
 
         htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
-        compute(d_spad, s0_spad, src1_plane, i01, ne11, &bctx->src1_dim1_div, nb11,
+        compute(d_spad, s0_spad, src1_plane, i01, ne11, &bctx->src1_dim1_div, vtcm_row_stride,
                 current_block_size, bctx->dst_row_size_aligned, bctx->src0_row_size_aligned, ne00);
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
 
@@ -750,6 +648,12 @@ static void binary_thread_element_repeat(unsigned int nth, unsigned int ith, voi
     uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_thread);
     size_t src0_spad_half    = layout->src0_spad_half_size;
     size_t dst_spad_half     = layout->dst_spad_half_size;
+    const uint8_t * vtcm_src1 = VTCM_LAYOUT_PTR(const uint8_t, bctx->vtcm_base, layout->off_src1);
+
+    const uint32_t elem_size = (src0->type == HTP_TYPE_F32) ? sizeof(float) : sizeof(_Float16);
+    const size_t vtcm_row_stride   = (size_t) ne10 * elem_size;
+    const size_t vtcm_plane_stride = (size_t) ne11 * vtcm_row_stride;
+    const size_t vtcm_vol_stride   = (size_t) ne12 * vtcm_plane_stride;
 
     FARF(HIGH, "binary-repeat: %d/%d (%u:%u) row-size %u (%u)",
          ith, nth, start_row, end_row, nb01, bctx->dst_row_size_aligned);
@@ -792,10 +696,10 @@ static void binary_thread_element_repeat(unsigned int nth, unsigned int ith, voi
 
         uint32_t i13 = fastmodulo(i03, ne13, &bctx->src1_dim3_div);
         uint32_t i12 = fastmodulo(i02, ne12, &bctx->src1_dim2_div);
-        const uint8_t * src1_plane = (const uint8_t *)(uintptr_t) src1->data + i13 * nb13 + i12 * nb12;
+        const uint8_t * src1_plane = vtcm_src1 + i13 * vtcm_vol_stride + i12 * vtcm_plane_stride;
 
         htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
-        compute(d_spad, s0_spad, src1_plane, i01, ne11, &bctx->src1_dim1_div, nb11,
+        compute(d_spad, s0_spad, src1_plane, i01, ne11, &bctx->src1_dim1_div, vtcm_row_stride,
                 current_block_size, bctx->dst_row_size_aligned, bctx->src0_row_size_aligned, ne00, ne10);
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
 
@@ -826,7 +730,6 @@ static void binary_thread_add_id_f32(unsigned int nth, unsigned int ith, void *
 
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * src1 = octx->src[1];
-    const struct htp_tensor * src2 = octx->src[2];
     const struct htp_tensor * dst  = octx->dst;
 
     const uint32_t ne00 = src0->ne[0];
@@ -852,6 +755,7 @@ static void binary_thread_add_id_f32(unsigned int nth, unsigned int ith, void *
     uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_thread);
     uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, bctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_thread);
     const uint8_t * vtcm_src1 = VTCM_LAYOUT_PTR(const uint8_t, bctx->vtcm_base, layout->off_src1);
+    const int32_t * vtcm_src2 = VTCM_LAYOUT_PTR(const int32_t, bctx->vtcm_base, layout->off_src2);
     size_t src0_spad_half    = layout->src0_spad_half_size;
     size_t dst_spad_half     = layout->dst_spad_half_size;
 
@@ -892,8 +796,8 @@ static void binary_thread_add_id_f32(unsigned int nth, unsigned int ith, void *
 
         htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
         compute_add_id_t compute = (compute_add_id_t) bctx->compute;
-        compute(d_spad, s0_spad, vtcm_src1, (const char *)(uintptr_t)src2->data,
-                i01, i02, src2->nb[0], src2->nb[1], src1_stride,
+        compute(d_spad, s0_spad, vtcm_src1, vtcm_src2,
+                i01, i02, ne01, src1_stride,
                 current_block_size, bctx->dst_row_size_aligned, bctx->src0_row_size_aligned, ne00);
         htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);
 
@@ -1234,19 +1138,6 @@ static int execute_op_binary(struct htp_ops_context * octx) {
     const size_t src1_row_size_aligned = kparams->src1_row_size_aligned;
     const size_t dst_row_size_aligned  = kparams->dst_row_size_aligned;
 
-    if (htp_tensor_is_extended(src1)) {
-        if (kparams->kernel_type != HTP_BINARY_KERNEL_SAME_SHAPE &&
-            kparams->kernel_type != HTP_BINARY_KERNEL_ROW_BCAST &&
-            kparams->kernel_type != HTP_BINARY_KERNEL_SCALAR_DMA &&
-            kparams->kernel_type != HTP_BINARY_KERNEL_ADD_ID) {
-            return HTP_STATUS_NO_SUPPORT;
-        }
-    }
-
-    if (octx->op == HTP_OP_ADD_ID && htp_tensor_is_extended(octx->src[2])) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
     struct htp_binary_context bctx;
     bctx.vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
     htp_binary_vtcm_layout_build(&bctx.vtcm_layout, kparams, octx->ctx->vtcm_size);
@@ -1259,12 +1150,38 @@ static int execute_op_binary(struct htp_ops_context * octx) {
     uint8_t * vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src1);
     if (kparams->kernel_type == HTP_BINARY_KERNEL_ROW_BCAST) {
         dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), bctx.vtcm_layout.src1_size, 0, src1->ne[0] * elem_size, 1);
-    } else if (kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR_DMA) {
-        dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), bctx.vtcm_layout.src1_size, 0, src1->ne[1] * elem_size, 1);
+    } else if (kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR ||
+               kparams->kernel_type == HTP_BINARY_KERNEL_REPEAT ||
+               kparams->kernel_type == HTP_BINARY_KERNEL_COMPLEX) {
+        if (!kparams->src1_strided) {
+            const size_t src1_bytes = (size_t) src1->ne[0] * src1->ne[1] * src1->ne[2] * src1->ne[3] * elem_size;
+            dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), bctx.vtcm_layout.src1_size, 0, src1_bytes, 1);
+        } else {
+            const uint32_t ne11 = src1->ne[1];
+            const uint32_t ne12 = src1->ne[2];
+            const uint32_t ne13 = src1->ne[3];
+            const size_t row_bytes = src1->ne[0] * elem_size;
+            for (uint32_t i13 = 0; i13 < ne13; i13++) {
+                for (uint32_t i12 = 0; i12 < ne12; i12++) {
+                    dma_addr_t src_addr = src1->data + i13 * src1->nb[3] + i12 * src1->nb[2];
+                    uint8_t * dst_ptr = vtcm_src1 + (i13 * ne12 * ne11 + i12 * ne11) * row_bytes;
+                    while (!dma_queue_push(dma_q, dma_make_data(dst_ptr, src_addr),
+                                           row_bytes, src1->nb[1], row_bytes, ne11)) {
+                        dma_queue_pop(dma_q);
+                    }
+                }
+            }
+        }
     } else if (kparams->kernel_type == HTP_BINARY_KERNEL_ADD_ID) {
+        const struct htp_tensor * src2 = octx->src[2];
+        uint8_t * vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, bctx.vtcm_base, bctx.vtcm_layout.off_src2);
         dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data),
                        kparams->src1_row_size_aligned, src1->nb[1],
                        src1->ne[0] * elem_size, src1->ne[1]);
+        const size_t src2_row_bytes = src2->ne[0] * sizeof(int32_t);
+        dma_queue_push(dma_q, dma_make_data(vtcm_src2, src2->data),
+                       src2_row_bytes, src2->nb[1],
+                       src2_row_bytes, src2->ne[1]);
     }
 
     bctx.octx                  = octx;
@@ -1285,15 +1202,8 @@ static int execute_op_binary(struct htp_ops_context * octx) {
     bctx.src1_dim2_div  = init_fastdiv_values(src1->ne[2]);
     bctx.src1_dim3_div  = init_fastdiv_values(src1->ne[3]);
 
-    bool src0_contig_dim1 = (src0->nb[2] == src0->ne[1] * src0->nb[1]);
-    bool dst_contig_dim1  = (dst->nb[2]  == src0->ne[1] * dst->nb[1]);
-
-    bool src0_contig_dim2 = (src0->nb[3] == src0->ne[2] * src0->nb[2]);
-    bool dst_contig_dim2  = (dst->nb[3]  == src0->ne[2] * dst->nb[2]);
-
-    bctx.split_at_ne01 = (octx->op == HTP_OP_ADD_ID) ||
-        ((src0->ne[2] > 1) && ((src1->ne[1] > 1) || (src1->ne[2] > 1) || !src0_contig_dim1 || !dst_contig_dim1));
-    bctx.split_at_ne02 = (src0->ne[3] > 1) && ((src1->ne[2] > 1) || (src1->ne[3] > 1) || !src0_contig_dim2 || !dst_contig_dim2);
+    bctx.split_at_ne01  = kparams->split_at_ne01;
+    bctx.split_at_ne02  = kparams->split_at_ne02;
 
     worker_callback_t worker_func = NULL;
     void * compute_func = NULL;
@@ -1339,26 +1249,6 @@ static int execute_op_binary(struct htp_ops_context * octx) {
                 }
             }
             break;
-        case HTP_BINARY_KERNEL_SCALAR_DMA:
-            worker_func = binary_thread_scalar_dma;
-            if (src0_type == HTP_TYPE_F32) {
-                switch (octx->op) {
-                    case HTP_OP_ADD: compute_func = compute_scalar_dma_add_f32; break;
-                    case HTP_OP_SUB: compute_func = compute_scalar_dma_sub_f32; break;
-                    case HTP_OP_MUL: compute_func = compute_scalar_dma_mul_f32; break;
-                    case HTP_OP_DIV: compute_func = compute_scalar_dma_div_f32; break;
-                    default: break;
-                }
-            } else if (src0_type == HTP_TYPE_F16) {
-                switch (octx->op) {
-                    case HTP_OP_ADD: compute_func = compute_scalar_dma_add_f16; break;
-                    case HTP_OP_SUB: compute_func = compute_scalar_dma_sub_f16; break;
-                    case HTP_OP_MUL: compute_func = compute_scalar_dma_mul_f16; break;
-                    case HTP_OP_DIV: compute_func = compute_scalar_dma_div_f16; break;
-                    default: break;
-                }
-            }
-            break;
         case HTP_BINARY_KERNEL_SCALAR:
             worker_func = binary_thread_scalar;
             if (src0_type == HTP_TYPE_F32) {
@@ -1435,9 +1325,11 @@ static int execute_op_binary(struct htp_ops_context * octx) {
     bctx.compute = compute_func;
 
     if (kparams->kernel_type == HTP_BINARY_KERNEL_ROW_BCAST ||
-        kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR_DMA ||
+        kparams->kernel_type == HTP_BINARY_KERNEL_SCALAR ||
+        kparams->kernel_type == HTP_BINARY_KERNEL_REPEAT ||
+        kparams->kernel_type == HTP_BINARY_KERNEL_COMPLEX ||
         kparams->kernel_type == HTP_BINARY_KERNEL_ADD_ID) {
-        dma_queue_pop(dma_q);
+        dma_queue_flush(dma_q);
     }
 
     work_queue_run(octx->ctx->work_queue, worker_func, &bctx, n_threads);
@@ -1446,13 +1338,6 @@ static int execute_op_binary(struct htp_ops_context * octx) {
 }
 
 int op_binary(struct htp_ops_context * octx) {
-
-    // Does not support permutations of src1
-    const struct htp_tensor * src1 = octx->src[1];
-    if (src1->nb[1] < src1->nb[0]) {
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
     const uint32_t src0_type = octx->src[0]->type;
     if ((src0_type == HTP_TYPE_F32) || (src0_type == HTP_TYPE_F16)) {
         return execute_op_binary(octx);

ggml/src/ggml-hexagon/htp/binary-ops.h

@@ -12,7 +12,6 @@
 enum htp_binary_kernel_type {
     HTP_BINARY_KERNEL_SAME_SHAPE = 0,
     HTP_BINARY_KERNEL_ROW_BCAST,
-    HTP_BINARY_KERNEL_SCALAR_DMA,
     HTP_BINARY_KERNEL_SCALAR,
     HTP_BINARY_KERNEL_ADD_ID,
     HTP_BINARY_KERNEL_COMPLEX,
@@ -30,15 +29,41 @@ struct htp_binary_kernel_params {
     uint32_t dst_row_size_aligned;
 
     uint32_t src1_size;
+    uint32_t src2_size;
     uint32_t vtcm_size;
 
     uint32_t chunk_size;
     uint32_t chunk_bytes;
     uint32_t is_scalar;
+
+    uint8_t  split_at_ne01;
+    uint8_t  split_at_ne02;
+    uint8_t  src1_strided;
+    uint8_t  pad;
 };
 
 #if defined(__cplusplus)
 static_assert(sizeof(struct htp_binary_kernel_params) <= 128, "htp_binary_kernel_params is too large for kernel_params blob");
+
+struct ggml_tensor;
+static inline bool htp_binary_strides_safe(
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * src1,
+    const struct ggml_tensor * dst,
+    size_t elem_size
+) {
+    if ((src0->ne[0] != 1 && src0->nb[0] != elem_size) ||
+        (src1->ne[0] != 1 && src1->nb[0] != elem_size) ||
+        (dst->ne[0]  != 1 && dst->nb[0]  != elem_size)) {
+        return false;
+    }
+    if ((src0->ne[1] > 1 && src0->nb[1] < src0->ne[0] * elem_size) ||
+        (src1->ne[1] > 1 && src1->nb[1] < src1->ne[0] * elem_size) ||
+        (dst->ne[1]  > 1 && dst->nb[1]  < dst->ne[0]  * elem_size)) {
+        return false;
+    }
+    return true;
+}
 #else
 _Static_assert(sizeof(struct htp_binary_kernel_params) <= 128, "htp_binary_kernel_params is too large for kernel_params blob");
 #endif
@@ -47,6 +72,7 @@ struct htp_binary_vtcm_layout {
     size_t total_bytes;
     size_t off_src0;
     size_t off_src1;
+    size_t off_src2;
     size_t off_dst;
 
     size_t src0_bytes_per_thread;
@@ -58,6 +84,7 @@ struct htp_binary_vtcm_layout {
     size_t dst_spad_half_size;
 
     size_t src1_size;
+    size_t src2_size;
     uint32_t rows_per_buffer;
 };
 
@@ -111,17 +138,19 @@ static inline void htp_binary_vtcm_layout_build(
         ? 2 * (kparams->src0_row_size_aligned + kparams->src1_row_size_aligned + kparams->dst_row_size_aligned)
         : 2 * (kparams->src0_row_size_aligned + kparams->dst_row_size_aligned);
 
-    if (spad_row_total == 0 || vtcm_size < kparams->src1_size) {
+    const size_t static_vtcm_size = kparams->src1_size + kparams->src2_size;
+    if (spad_row_total == 0 || vtcm_size < static_vtcm_size) {
         return;
     }
 
-    const size_t rows_per_buffer = (vtcm_size - kparams->src1_size) / (n_threads * spad_row_total);
+    const size_t rows_per_buffer = (vtcm_size - static_vtcm_size) / (n_threads * spad_row_total);
     if (rows_per_buffer == 0) {
         return;
     }
 
     L->rows_per_buffer = (uint32_t) rows_per_buffer;
     L->src1_size = kparams->src1_size;
+    L->src2_size = kparams->src2_size;
 
     L->src0_bytes_per_thread = rows_per_buffer * 2 * kparams->src0_row_size_aligned;
     L->dst_bytes_per_thread  = rows_per_buffer * 2 * kparams->dst_row_size_aligned;
@@ -137,11 +166,13 @@ static inline void htp_binary_vtcm_layout_build(
     const size_t src1_total = (kparams->src1_size > 0)
         ? kparams->src1_size
         : n_threads * L->src1_bytes_per_thread;
+    const size_t src2_total = kparams->src2_size;
     const size_t dst_total  = n_threads * L->dst_bytes_per_thread;
 
     size_t off = 0;
     VTCM_LAYOUT_ALLOC(off, off_src0, src0_total);
     VTCM_LAYOUT_ALLOC(off, off_src1, src1_total);
+    VTCM_LAYOUT_ALLOC(off, off_src2, src2_total);
     VTCM_LAYOUT_ALLOC(off, off_dst,  dst_total);
 
     L->total_bytes = off;

ggml/src/ggml-vulkan/ggml-vulkan-buffers.cpp

@@ -1,5 +1,27 @@
 #include "ggml-vulkan-common.h"
 
+#include <atomic>
+#include <cstdint>
+#include <cstdlib>
+#include <mutex>
+#include <shared_mutex>
+
+static void * portable_aligned_malloc(size_t align, size_t size) {
+#ifdef _WIN32
+    return _aligned_malloc(size, align);
+#else
+    return std::aligned_alloc(align, size);
+#endif
+}
+
+static void portable_aligned_free(void * ptr) {
+#ifdef _WIN32
+    _aligned_free(ptr);
+#else
+    std::free(ptr);
+#endif
+}
+
 ggml_backend_buffer_type_i ggml_backend_vk_buffer_type_interface = {
     /* .get_name            = */ ggml_backend_vk_buffer_type_name,
     /* .alloc_buffer        = */ ggml_backend_vk_buffer_type_alloc_buffer,
@@ -243,67 +265,125 @@ void ggml_vk_destroy_buffer(vk_buffer& buf) {
     buf.reset();
 }
 
-void * ggml_vk_host_malloc(vk_device& device, size_t size) {
-    VK_LOG_MEMORY("ggml_vk_host_malloc(" << size << ")");
-    vk_buffer buf = ggml_vk_create_buffer(device, size,
-        {vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached,
-         vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent});
+// caller must hold host_allocations_mutex
+static size_t ggml_vk_host_allocation_find(const void * ptr) {
+    for (size_t i = 0; i < vk_instance.host_allocations.size(); i++) {
+        const uint8_t * base = (const uint8_t *) vk_instance.host_allocations[i].base;
+        if (ptr >= base && ptr < base + vk_instance.host_allocations[i].size) {
+            return i;
+        }
+    }
+    return SIZE_MAX;
+}
 
-    if(!(buf->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible)) {
-        fprintf(stderr, "WARNING: failed to allocate %.2f MB of pinned memory\n",
-            size/1024.0/1024.0);
-        device->device.freeMemory(buf->device_memory);
-        device->device.destroyBuffer(buf->buffer);
+void * ggml_vk_host_malloc(size_t size) {
+    VK_LOG_MEMORY("ggml_vk_host_malloc(" << size << ")");
+
+    ggml_vk_instance_init();
+
+    if (vk_instance.multi_device) {
+        const size_t align      = std::max<size_t>(vk_instance.host_import_alignment, TENSOR_ALIGNMENT);
+        const size_t alloc_size = (size + align - 1) / align * align;
+
+        void * ptr = portable_aligned_malloc(align, alloc_size);
+        if (ptr == nullptr) {
+            return nullptr;
+        }
+        vk_host_allocation entry;
+        entry.base = ptr;
+        entry.size = alloc_size;
+
+        std::unique_lock<std::shared_mutex> guard(vk_instance.host_allocations_mutex);
+        vk_instance.host_allocations.push_back(std::move(entry));
+        return ptr;
+    }
+
+    vk_device dev0 = ggml_vk_get_device(0);
+    vk_buffer buf;
+    try {
+        buf = ggml_vk_create_buffer(dev0, size,
+            {vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached,
+             vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent});
+    } catch (vk::SystemError&) {
         return nullptr;
     }
 
-    std::lock_guard<std::shared_mutex> guard(device->pinned_memory_mutex);
-    device->pinned_memory.push_back(std::make_tuple(buf->ptr, size, buf));
+    if (!(buf->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible)) {
+        fprintf(stderr, "WARNING: failed to allocate %.2f MB of pinned memory\n", size / 1024.0 / 1024.0);
+        ggml_vk_destroy_buffer(buf);
+        return nullptr;
+    }
 
+    vk_host_allocation entry;
+    entry.base             = buf->ptr;
+    entry.size             = size;
+    entry.views[dev0->idx] = buf;
+
+    std::unique_lock<std::shared_mutex> guard(vk_instance.host_allocations_mutex);
+    vk_instance.host_allocations.push_back(std::move(entry));
     return buf->ptr;
 }
 
-void ggml_vk_host_free(vk_device& device, void* ptr) {
+void ggml_vk_host_free(void * ptr) {
     if (ptr == nullptr) {
         return;
     }
     VK_LOG_MEMORY("ggml_vk_host_free(" << ptr << ")");
-    std::lock_guard<std::shared_mutex> guard(device->pinned_memory_mutex);
 
-    vk_buffer buf;
-    size_t index;
-    for (size_t i = 0; i < device->pinned_memory.size(); i++) {
-        const uint8_t* addr = (const uint8_t*) std::get<0>(device->pinned_memory[i]);
-        const uint8_t* endr = addr + std::get<1>(device->pinned_memory[i]);
-        if (ptr >= addr && ptr < endr) {
-            buf = std::get<2>(device->pinned_memory[i]);
-            index = i;
-            break;
+    vk_host_allocation entry;
+    {
+        std::unique_lock<std::shared_mutex> guard(vk_instance.host_allocations_mutex);
+        const size_t index = ggml_vk_host_allocation_find(ptr);
+        if (index == SIZE_MAX) {
+            fprintf(stderr, "WARNING: failed to free pinned memory: memory not in map\n");
+            return;
         }
-    }
-    if (buf == nullptr) {
-        fprintf(stderr, "WARNING: failed to free pinned memory: memory not in map\n");
-        return;
+        entry = std::move(vk_instance.host_allocations[index]);
+        vk_instance.host_allocations.erase(vk_instance.host_allocations.begin() + index);
     }
 
-    ggml_vk_destroy_buffer(buf);
-
-    device->pinned_memory.erase(device->pinned_memory.begin() + index);
+    for (auto & view : entry.views) {
+        ggml_vk_destroy_buffer(view);
+    }
+    if (vk_instance.multi_device) {
+        portable_aligned_free(entry.base);
+    }
 }
 
 void ggml_vk_host_get(const vk_device& device, const void * ptr, vk_buffer& buf, size_t& buf_offset) {
-    std::shared_lock<std::shared_mutex> guard(device->pinned_memory_mutex);
     buf = nullptr;
     buf_offset = 0;
-    for (size_t i = 0; i < device->pinned_memory.size(); i++) {
-        const uint8_t* addr = (const uint8_t*) std::get<0>(device->pinned_memory[i]);
-        const uint8_t* endr = addr + std::get<1>(device->pinned_memory[i]);
-        if (ptr >= addr && ptr < endr) {
-            buf = std::get<2>(device->pinned_memory[i]);
-            buf_offset = ((const uint8_t *)ptr) - addr;
-            break;
+
+    std::unique_lock<std::shared_mutex> guard(vk_instance.host_allocations_mutex);
+
+    const size_t index = ggml_vk_host_allocation_find(ptr);
+    if (index == SIZE_MAX) {
+        return;
+    }
+    vk_host_allocation & entry = vk_instance.host_allocations[index];
+
+    if (entry.views[device->idx] == nullptr) {
+        if (!vk_instance.multi_device) {
+            return;
+        }
+        if (!device->external_memory_host) {
+            static std::atomic<uint32_t> warned_mask{0};
+            const uint32_t bit = 1u << device->idx;
+            if (!(warned_mask.fetch_or(bit) & bit)) {
+                GGML_LOG_WARN("ggml_vulkan: %s does not support VK_EXT_external_memory_host; host memory "
+                    "will be staged, not pinned, on this device\n", device->name.c_str());
+            }
+            return;
+        }
+        // const only propagates from the const context; the pointee is mutable
+        entry.views[device->idx] = ggml_vk_buffer_from_host_ptr(const_cast<vk_device&>(device), entry.base, entry.size);
+        if (entry.views[device->idx] == nullptr) {
+            return;
         }
     }
+
+    buf = entry.views[device->idx];
+    buf_offset = (const uint8_t *) ptr - (const uint8_t *) entry.base;
 }
 
 void ggml_vk_ensure_sync_staging_buffer(vk_device& device, size_t size) {

ggml/src/ggml-vulkan/ggml-vulkan-common.h

@@ -45,8 +45,8 @@ bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_
 vk_buffer ggml_vk_create_buffer_check(vk_device& device, size_t size, vk::MemoryPropertyFlags req_flags, vk::MemoryPropertyFlags fallback_flags = vk::MemoryPropertyFlags(0));
 vk_buffer ggml_vk_create_buffer_device(vk_device& device, size_t size);
 void ggml_vk_destroy_buffer(vk_buffer& buf);
-void * ggml_vk_host_malloc(vk_device& device, size_t size);
-void ggml_vk_host_free(vk_device& device, void* ptr);
+void * ggml_vk_host_malloc(size_t size);
+void ggml_vk_host_free(void * ptr);
 void ggml_vk_host_get(const vk_device& device, const void * ptr, vk_buffer& buf, size_t& buf_offset);
 void ggml_vk_ensure_sync_staging_buffer(vk_device& device, size_t size);
 void ggml_vk_ensure_sync_staging_buffer(ggml_backend_vk_context * ctx, size_t size);
@@ -94,7 +94,7 @@ void deferred_memset(void * dst, uint32_t val, size_t size, std::vector<vk_stagi
 
 // matmul
 vk_pipeline ggml_vk_get_to_fp16(ggml_backend_vk_context * ctx, ggml_type type);
-void ggml_vk_matmul(ggml_backend_vk_context * ctx, vk_context& subctx, vk_pipeline& pipeline, vk_subbuffer&& a, vk_subbuffer&& b, vk_subbuffer&& d, vk_subbuffer&& split_k_buffer, uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d, uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d, uint32_t split_k, uint32_t batch, uint32_t ne02, uint32_t ne12, uint32_t broadcast2, uint32_t broadcast3, uint32_t padded_n);
+void ggml_vk_matmul(ggml_backend_vk_context * ctx, vk_context& subctx, vk_pipeline& pipeline, vk_subbuffer&& a, vk_subbuffer&& b, vk_subbuffer&& d, vk_subbuffer&& split_k_buffer, uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d, uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d, uint32_t split_k, uint32_t batch, uint32_t ne02, uint32_t ne12, uint32_t broadcast2, uint32_t broadcast3, uint32_t padded_n, uint64_t b_col_bytes, uint64_t d_col_bytes);
 bool ggml_vk_dim01_contiguous(const ggml_tensor * tensor);
 vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const ggml_tensor * src, const ggml_tensor * dst, ggml_type to);
 vk_pipeline ggml_vk_get_quantize_pipeline(ggml_backend_vk_context * ctx, ggml_type type);

ggml/src/ggml-vulkan/ggml-vulkan-debug.cpp

@@ -173,7 +173,8 @@ void ggml_vk_test_matmul(ggml_backend_vk_context * ctx, size_t m, size_t n, size
             ctx, subctx, p, ggml_vk_subbuffer(ctx, d_X), ggml_vk_subbuffer(ctx, d_Y), ggml_vk_subbuffer(ctx, d_D), ggml_vk_subbuffer(ctx, ctx->prealloc_split_k),
             m, n, k,
             k, k, m, k*m, k*n, m*n,
-            split_k, batch, batch, batch, 1, 1, n
+            split_k, batch, batch, batch, 1, 1, n,
+            k * sizeof(Y_TYPE), m * sizeof(float)
         );
     }
     ggml_vk_ctx_end(subctx);
@@ -640,7 +641,8 @@ void ggml_vk_test_dequant_matmul(ggml_backend_vk_context * ctx, size_t m, size_t
                 ctx, subctx, p, { qx_buf, 0, qx_sz }, { qy_buf, 0, qy_sz }, { d_buf, 0, d_sz }, { ctx->prealloc_split_k, 0, ctx->prealloc_size_split_k },
                 m, n, k,
                 k, k, m, k*m, k*n, m*n,
-                split_k, batch, batch, batch, 1, 1, n
+                split_k, batch, batch, batch, 1, 1, n,
+                (uint64_t)(k / ggml_blck_size(GGML_TYPE_Q8_1)) * ggml_type_size(GGML_TYPE_Q8_1), m * sizeof(float)
             );
         }
     } else {
@@ -649,7 +651,8 @@ void ggml_vk_test_dequant_matmul(ggml_backend_vk_context * ctx, size_t m, size_t
                 ctx, subctx, p, { qx_buf, 0, qx_sz }, { y_buf, 0, y_sz }, { d_buf, 0, d_sz }, { ctx->prealloc_split_k, 0, ctx->prealloc_size_split_k },
                 m, n, k,
                 k, k, m, k*m, k*n, m*n,
-                split_k, batch, batch, batch, 1, 1, n
+                split_k, batch, batch, batch, 1, 1, n,
+                k * sizeof(float), m * sizeof(float)
             );
         }
     }

ggml/src/ggml-vulkan/ggml-vulkan-types.h

@@ -48,43 +48,24 @@ typedef struct VkPhysicalDeviceCooperativeMatrixDecodeVectorFeaturesNV {
 #endif
 
 #include <algorithm>
-
 #include <cmath>
-
 #include <functional>
-
 #include <iomanip>
-
 #include <iostream>
-
 #include <tuple>
-
 #include <vector>
-
 #include <deque>
-
 #include <sstream>
-
 #include <utility>
-
 #include <memory>
-
 #include <limits>
-
 #include <map>
-
 #include <set>
-
 #include <unordered_map>
-
 #include <shared_mutex>
-
 #include <mutex>
-
 #include <future>
-
 #include <condition_variable>
-
 #include <thread>
 
 #if defined(_MSC_VER)
@@ -110,9 +91,7 @@ typedef struct VkPhysicalDeviceCooperativeMatrixDecodeVectorFeaturesNV {
 #endif
 
 #include "ggml-impl.h"
-
 #include "ggml-backend-impl.h"
-
 #include "ggml-vulkan-shaders.hpp"
 
 #if !defined(VK_KHR_shader_bfloat16)
@@ -680,7 +659,6 @@ class vk_memory_logger;
 struct vk_device_struct {
     std::recursive_mutex mutex;
     std::mutex queue_submit_mutex;
-    mutable std::shared_mutex pinned_memory_mutex;
 
     // Guards compile_pending, all_pipelines, and the dynamic pipeline maps
     // (flash_attn, fa_mask_opt, solve_tri, conv2d, etc). The actual compile
@@ -1019,8 +997,6 @@ struct vk_device_struct {
 
     std::vector<vk_pipeline_ref> all_pipelines;
 
-    std::vector<std::tuple<void*, size_t, vk_buffer>> pinned_memory;
-
     vk::Fence fence;
     vk_buffer sync_staging;
 
@@ -1333,6 +1309,12 @@ struct ggml_backend_vk_buffer_context {
 
 };
 
+struct vk_host_allocation {
+    void *    base = nullptr;
+    size_t    size = 0;
+    vk_buffer views[GGML_VK_MAX_DEVICES];  // indexed by device->idx; null = not (yet) mapped here
+};
+
 struct vk_instance_t {
     vk::Instance instance;
 
@@ -1347,6 +1329,13 @@ struct vk_instance_t {
     std::vector<size_t> device_indices;
     std::vector<bool>   device_supports_membudget;
     vk_device devices[GGML_VK_MAX_DEVICES];
+
+    std::vector<vk_host_allocation> host_allocations;
+    std::shared_mutex               host_allocations_mutex;
+
+    size_t host_import_alignment = 0;
+
+    bool multi_device = false;
 };
 
 typedef void (*ggml_vk_func_t)(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst);

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -5411,22 +5411,37 @@ void ggml_vk_instance_init() {
     }
     GGML_LOG_DEBUG("ggml_vulkan: Found %zu Vulkan devices:\n", vk_instance.device_indices.size());
 
+    bool all_external_memory_host = vk_instance.device_indices.size() > 1;
     for (size_t i = 0; i < vk_instance.device_indices.size(); i++) {
         vk::PhysicalDevice vkdev = devices[vk_instance.device_indices[i]];
         std::vector<vk::ExtensionProperties> extensionprops = vkdev.enumerateDeviceExtensionProperties();
 
         bool membudget_supported = false;
+        bool external_memory_host_supported = false;
         for (const auto & ext : extensionprops) {
             if (strcmp(VK_EXT_MEMORY_BUDGET_EXTENSION_NAME, ext.extensionName) == 0) {
                 membudget_supported = true;
-                break;
+            } else if (strcmp("VK_EXT_external_memory_host", ext.extensionName) == 0) {
+                external_memory_host_supported = true;
             }
         }
 
         vk_instance.device_supports_membudget.push_back(membudget_supported);
 
+        if (external_memory_host_supported) {
+            vk::PhysicalDeviceProperties2 props2;
+            vk::PhysicalDeviceExternalMemoryHostPropertiesEXT ext_props;
+            props2.pNext = &ext_props;
+            vkdev.getProperties2(&props2);
+            vk_instance.host_import_alignment = std::max<size_t>(vk_instance.host_import_alignment, ext_props.minImportedHostPointerAlignment);
+        } else {
+            all_external_memory_host = false;
+        }
+
         ggml_vk_print_gpu_info(i);
     }
+
+    vk_instance.multi_device = all_external_memory_host;
 }
 
 void ggml_vk_init(ggml_backend_vk_context * ctx, size_t idx) {
@@ -5907,18 +5922,34 @@ void ggml_vk_matmul(
         uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d,
         uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d,
         uint32_t split_k, uint32_t batch, uint32_t ne02, uint32_t ne12, uint32_t broadcast2, uint32_t broadcast3,
-        uint32_t padded_n) {
+        uint32_t padded_n, uint64_t b_col_bytes, uint64_t d_col_bytes) {
         VK_LOG_DEBUG("ggml_vk_matmul(a: (" << a.buffer->buffer << ", " << a.offset << ", " << a.size << "), b: (" << b.buffer->buffer << ", " << b.offset << ", " << b.size << "), d: (" << d.buffer->buffer << ", " << d.offset << ", " << d.size << "), split_k: (" << (split_k_buffer.buffer != nullptr ? split_k_buffer.buffer->buffer : VK_NULL_HANDLE) << ", " << split_k_buffer.offset << ", " << split_k_buffer.size << "), m: " << m << ", n: " << n << ", k: " << k << ", stride_a: " << stride_a << ", stride_b: " << stride_b << ", stride_d: " << stride_d << ", batch_stride_a: " << batch_stride_a << ", batch_stride_b: " << batch_stride_b << ", batch_stride_d: " << batch_stride_d << ", split_k: " << split_k << ", batch: " << batch << ", ne02: " << ne02 << ", ne12: " << ne12 << ", broadcast2: " << broadcast2 << ", broadcast3: " << broadcast3 << ", padded_n: " << padded_n << ")");
     if (split_k == 1) {
-        ggml_pipeline_request_descriptor_sets(ctx, pipeline, CEIL_DIV(batch, ctx->device->properties.limits.maxComputeWorkGroupCount[2]));
+        // a small tile can still produce more workgroups than the device allows on the n axis:
+        // split the dispatch into n-chunks, shifting the B and D column windows
+        const uint64_t max_n = (uint64_t)ctx->device->properties.limits.maxComputeWorkGroupCount[1] * pipeline->wg_denoms[1];
+        const uint64_t n_chunks = CEIL_DIV(n, max_n);
+        ggml_pipeline_request_descriptor_sets(ctx, pipeline, n_chunks * CEIL_DIV(batch, ctx->device->properties.limits.maxComputeWorkGroupCount[2]));
 
-        uint32_t base_work_group_z = 0;
-        while (base_work_group_z < batch) {
-            uint32_t groups_z = std::min(batch - base_work_group_z, ctx->device->properties.limits.maxComputeWorkGroupCount[2]);
+        for (uint64_t n_base = 0; n_base < n; n_base += max_n) {
+            const uint32_t n_chunk = (uint32_t) std::min(n - n_base, max_n);
+            GGML_ASSERT(n_base * b_col_bytes < b.size && n_base * d_col_bytes < d.size);
+            GGML_ASSERT(b.size != VK_WHOLE_SIZE && d.size != VK_WHOLE_SIZE);
+            vk_subbuffer b_chunk = b;
+            vk_subbuffer d_chunk = d;
+            b_chunk.offset += n_base * b_col_bytes;
+            b_chunk.size   -= n_base * b_col_bytes;
+            d_chunk.offset += n_base * d_col_bytes;
+            d_chunk.size   -= n_base * d_col_bytes;
 
-            const vk_mat_mat_push_constants pc = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d, base_work_group_z, batch, k, ne02, ne12, broadcast2, broadcast3, padded_n };
-            ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d }, pc, { m, n, groups_z });
-            base_work_group_z += groups_z;
+            uint32_t base_work_group_z = 0;
+            while (base_work_group_z < batch) {
+                uint32_t groups_z = std::min(batch - base_work_group_z, ctx->device->properties.limits.maxComputeWorkGroupCount[2]);
+
+                const vk_mat_mat_push_constants pc = { m, n_chunk, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d, base_work_group_z, batch, k, ne02, ne12, broadcast2, broadcast3, uint32_t(padded_n - n_base) };
+                ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b_chunk, d_chunk }, pc, { m, n_chunk, groups_z });
+                base_work_group_z += groups_z;
+            }
         }
         return;
     }
@@ -6539,6 +6570,11 @@ static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& sub
         stride_batch_y = src1->nb[0] / ggml_type_size(src1->type);
     }
 
+    // column stride in bytes of the B buffer: raw src1 keeps its own stride, staging buffers are tightly packed
+    const uint64_t b_col_bytes = (y_non_contig || qy_needs_dequant || quantize_y)
+        ? (quantize_y ? (uint64_t)(ne10 / ggml_blck_size(GGML_TYPE_Q8_1)) * ggml_type_size(GGML_TYPE_Q8_1) : (uint64_t)ne10 * sizeof(ggml_fp16_t))
+        : (uint64_t)src1->nb[1];
+
     // compute
     ggml_vk_matmul(
         ctx, subctx, pipeline,
@@ -6546,7 +6582,8 @@ static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& sub
         ggml_vk_subbuffer(ctx, d_D, d_buf_offset), { ctx->prealloc_split_k, 0, d_sz * split_k },
         ne01, ne11, ne10,
         ne10, ne10, stride_d, stride_batch_x, stride_batch_y, stride_batch_d,
-        split_k, ne12*ne13, ne02, ne12, r2, r3, padded_n
+        split_k, ne12*ne13, ne02, ne12, r2, r3, padded_n,
+        b_col_bytes, (uint64_t)dst->nb[1]
     );  // NOLINT
 
     if (x_non_contig || qx_needs_dequant) {
@@ -13096,18 +13133,16 @@ static const char * ggml_backend_vk_host_buffer_type_name(ggml_backend_buffer_ty
 
 static void ggml_backend_vk_host_buffer_free_buffer(ggml_backend_buffer_t buffer) {
     VK_LOG_MEMORY("ggml_backend_vk_host_buffer_free_buffer()");
-    ggml_vk_host_free(vk_instance.devices[0], buffer->context);
+    ggml_vk_host_free(buffer->context);
 }
 
 static ggml_backend_buffer_t ggml_backend_vk_host_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) {
     VK_LOG_MEMORY("ggml_backend_vk_host_buffer_type_alloc_buffer(" << size << ")");
 
     size += 32;  // Behave like the CPU buffer type
-    void * ptr = nullptr;
-    try {
-        ptr = ggml_vk_host_malloc(vk_instance.devices[0], size);
-    } catch (vk::SystemError& e) {
-        GGML_LOG_WARN("ggml_vulkan: Failed to allocate pinned memory (%s)\n", e.what());
+    void * ptr = ggml_vk_host_malloc(size);
+    if (ptr == nullptr) {
+        GGML_LOG_WARN("ggml_vulkan: Failed to allocate pinned memory\n");
         // fallback to cpu buffer
         return ggml_backend_buft_alloc_buffer(ggml_backend_cpu_buffer_type(), size);
     }

ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp

@@ -201,9 +201,7 @@ void main() {
             const uint iqs = loadr_a;
 
             [[unroll]] for (uint k_step = 0; k_step < BK_STEP; k_step++) {
-                if (block + k_step * BK < end_k) {
-                    block_a_to_shmem(k_step * BM + buf_ib, ib + k_step, iqs);
-                }
+                block_a_to_shmem(k_step * BM + buf_ib, ib + k_step, iqs, block + k_step * BK < end_k);
             }
         }
         [[unroll]] for (uint l = 0; loadc_b + l < BN; l += loadstride_b) {

ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_funcs.glsl

@@ -7,18 +7,26 @@
 // Each iqs value maps to a 32-bit integer
 
 #if defined(DATA_A_Q2_0)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
-    const uint block_idx = ib / 2;
-    const uint byte_idx = (ib & 1u) * 8u + iqs;
-    const uint bits = uint(data_a[block_idx].qs[byte_idx]);
-    buf_a[buf_ib].qs[iqs] = pack32(i8vec4(
-        int8_t(bits & 3u),
-        int8_t((bits >> 2u) & 3u),
-        int8_t((bits >> 4u) & 3u),
-        int8_t(bits >> 6u)));
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+    if (is_in_bounds) {
+        const uint block_idx = ib / 2;
+        const uint byte_idx = (ib & 1u) * 8u + iqs;
+        const uint bits = uint(data_a[block_idx].qs[byte_idx]);
+        buf_a[buf_ib].qs[iqs] = pack32(i8vec4(
+            int8_t(bits & 3u),
+            int8_t((bits >> 2u) & 3u),
+            int8_t((bits >> 4u) & 3u),
+            int8_t(bits >> 6u)));
 
-    if (iqs == 0) {
-        buf_a[buf_ib].dm = FLOAT_TYPE(data_a[block_idx].d);
+        if (iqs == 0) {
+            buf_a[buf_ib].dm = FLOAT_TYPE(data_a[block_idx].d);
+        }
+    } else {
+        buf_a[buf_ib].qs[iqs] = 0;
+
+        if (iqs == 0) {
+            buf_a[buf_ib].dm = FLOAT_TYPE(0.0f);
+        }
     }
 }
 
@@ -43,21 +51,37 @@ ACC_TYPE mmq_dot_product(const uint ib_a) {
 #if defined(DATA_A_Q4_0) || defined(DATA_A_Q4_1)
 // 2-byte loads for Q4_0 blocks (18 bytes)
 // 4-byte loads for Q4_1 blocks (20 bytes)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+    if (is_in_bounds) {
+    #ifdef DATA_A_Q4_0
+        buf_a[buf_ib].qs[iqs] = pack32(u16vec2(data_a_packed16[ib].qs[iqs * 2],
+                                               data_a_packed16[ib].qs[iqs * 2 + 1]));
+
+        if (iqs == 0) {
+            buf_a[buf_ib].dm = FLOAT_TYPE(data_a_packed16[ib].d);
+        }
+    #else // DATA_A_Q4_1
+        buf_a[buf_ib].qs[iqs] = data_a_packed32[ib].qs[iqs];
+
+        if (iqs == 0) {
+            buf_a[buf_ib].dm = FLOAT_TYPEV2(data_a_packed32[ib].dm);
+        }
+    #endif
+    } else {
 #ifdef DATA_A_Q4_0
-    buf_a[buf_ib].qs[iqs] = pack32(u16vec2(data_a_packed16[ib].qs[iqs * 2],
-                                           data_a_packed16[ib].qs[iqs * 2 + 1]));
+        buf_a[buf_ib].qs[iqs] = 0;
 
-    if (iqs == 0) {
-        buf_a[buf_ib].dm = FLOAT_TYPE(data_a_packed16[ib].d);
-    }
+        if (iqs == 0) {
+            buf_a[buf_ib].dm = FLOAT_TYPE(0.0f);
+        }
 #else // DATA_A_Q4_1
-    buf_a[buf_ib].qs[iqs] = data_a_packed32[ib].qs[iqs];
+        buf_a[buf_ib].qs[iqs] = 0;
 
-    if (iqs == 0) {
-        buf_a[buf_ib].dm = FLOAT_TYPEV2(data_a_packed32[ib].dm);
-    }
+        if (iqs == 0) {
+            buf_a[buf_ib].dm = FLOAT_TYPEV2(0.0f);
+        }
 #endif
+    }
 }
 
 void block_a_to_registers(const uint reg_ib, const uint buf_ib) {
@@ -93,23 +117,41 @@ ACC_TYPE mmq_dot_product(const uint ib_a) {
 #if defined(DATA_A_Q5_0) || defined(DATA_A_Q5_1)
 // 2-byte loads for Q5_0 blocks (22 bytes)
 // 4-byte loads for Q5_1 blocks (24 bytes)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+    if (is_in_bounds) {
+    #ifdef DATA_A_Q5_0
+        buf_a[buf_ib].qs[iqs] = pack32(u16vec2(data_a_packed16[ib].qs[iqs * 2],
+                                               data_a_packed16[ib].qs[iqs * 2 + 1]));
+
+        if (iqs == 0) {
+            buf_a[buf_ib].dm = FLOAT_TYPE(data_a_packed16[ib].d);
+            buf_a[buf_ib].qh = pack32(u16vec2(data_a_packed16[ib].qh[0], data_a_packed16[ib].qh[1]));
+        }
+    #else // DATA_A_Q5_1
+        buf_a[buf_ib].qs[iqs] = data_a_packed32[ib].qs[iqs];
+
+        if (iqs == 0) {
+            buf_a[buf_ib].dm = FLOAT_TYPEV2(data_a_packed32[ib].dm);
+            buf_a[buf_ib].qh = data_a_packed32[ib].qh;
+        }
+    #endif
+    } else {
 #ifdef DATA_A_Q5_0
-    buf_a[buf_ib].qs[iqs] = pack32(u16vec2(data_a_packed16[ib].qs[iqs * 2],
-                                           data_a_packed16[ib].qs[iqs * 2 + 1]));
+        buf_a[buf_ib].qs[iqs] = 0;
 
-    if (iqs == 0) {
-        buf_a[buf_ib].dm = FLOAT_TYPE(data_a_packed16[ib].d);
-        buf_a[buf_ib].qh = pack32(u16vec2(data_a_packed16[ib].qh[0], data_a_packed16[ib].qh[1]));
-    }
+        if (iqs == 0) {
+            buf_a[buf_ib].dm = FLOAT_TYPE(0.0f);
+            buf_a[buf_ib].qh = 0;
+        }
 #else // DATA_A_Q5_1
-    buf_a[buf_ib].qs[iqs] = data_a_packed32[ib].qs[iqs];
+        buf_a[buf_ib].qs[iqs] = 0;
 
-    if (iqs == 0) {
-        buf_a[buf_ib].dm = FLOAT_TYPEV2(data_a_packed32[ib].dm);
-        buf_a[buf_ib].qh = data_a_packed32[ib].qh;
-    }
+        if (iqs == 0) {
+            buf_a[buf_ib].dm = FLOAT_TYPEV2(0.0f);
+            buf_a[buf_ib].qh = 0;
+        }
 #endif
+    }
 }
 
 void block_a_to_registers(const uint reg_ib, const uint buf_ib) {
@@ -148,12 +190,20 @@ ACC_TYPE mmq_dot_product(const uint ib_a) {
 
 #if defined(DATA_A_Q8_0)
 // 2-byte loads for Q8_0 blocks (34 bytes)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
-    buf_a[buf_ib].qs[iqs] = pack32(i16vec2(data_a_packed16[ib].qs[iqs * 2],
-                                           data_a_packed16[ib].qs[iqs * 2 + 1]));
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+    if (is_in_bounds) {
+        buf_a[buf_ib].qs[iqs] = pack32(i16vec2(data_a_packed16[ib].qs[iqs * 2],
+                                               data_a_packed16[ib].qs[iqs * 2 + 1]));
 
-    if (iqs == 0) {
-        buf_a[buf_ib].dm = FLOAT_TYPE(data_a_packed16[ib].d);
+        if (iqs == 0) {
+            buf_a[buf_ib].dm = FLOAT_TYPE(data_a_packed16[ib].d);
+        }
+    } else {
+        buf_a[buf_ib].qs[iqs] = 0;
+
+        if (iqs == 0) {
+            buf_a[buf_ib].dm = FLOAT_TYPE(0.0f);
+        }
     }
 }
 
@@ -180,20 +230,29 @@ ACC_TYPE mmq_dot_product(const uint ib_a) {
 
 #if defined(DATA_A_MXFP4)
 // 1-byte loads for mxfp4 blocks (17 bytes)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
-    const uint32_t qs = pack32(u8vec4(data_a[ib].qs[iqs * 4    ],
-                                      data_a[ib].qs[iqs * 4 + 1],
-                                      data_a[ib].qs[iqs * 4 + 2],
-                                      data_a[ib].qs[iqs * 4 + 3]));
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+    if (is_in_bounds) {
+        const uint32_t qs = pack32(u8vec4(data_a[ib].qs[iqs * 4    ],
+                                          data_a[ib].qs[iqs * 4 + 1],
+                                          data_a[ib].qs[iqs * 4 + 2],
+                                          data_a[ib].qs[iqs * 4 + 3]));
 
-    const u8vec4 i_a0 = unpack8( qs       & 0x0F0F0F0F);
-    const u8vec4 i_a1 = unpack8((qs >> 4) & 0x0F0F0F0F);
+        const u8vec4 i_a0 = unpack8( qs       & 0x0F0F0F0F);
+        const u8vec4 i_a1 = unpack8((qs >> 4) & 0x0F0F0F0F);
 
-    buf_a[buf_ib].qs[iqs    ] = pack32(i8vec4(kvalues_mxfp4[i_a0.x], kvalues_mxfp4[i_a0.y], kvalues_mxfp4[i_a0.z], kvalues_mxfp4[i_a0.w]));
-    buf_a[buf_ib].qs[iqs + 4] = pack32(i8vec4(kvalues_mxfp4[i_a1.x], kvalues_mxfp4[i_a1.y], kvalues_mxfp4[i_a1.z], kvalues_mxfp4[i_a1.w]));
+        buf_a[buf_ib].qs[iqs    ] = pack32(i8vec4(kvalues_mxfp4[i_a0.x], kvalues_mxfp4[i_a0.y], kvalues_mxfp4[i_a0.z], kvalues_mxfp4[i_a0.w]));
+        buf_a[buf_ib].qs[iqs + 4] = pack32(i8vec4(kvalues_mxfp4[i_a1.x], kvalues_mxfp4[i_a1.y], kvalues_mxfp4[i_a1.z], kvalues_mxfp4[i_a1.w]));
 
-    if (iqs == 0) {
-        buf_a[buf_ib].d = FLOAT_TYPE(e8m0_to_fp32(data_a[ib].e) * 0.5);
+        if (iqs == 0) {
+            buf_a[buf_ib].d = FLOAT_TYPE(e8m0_to_fp32(data_a[ib].e) * 0.5);
+        }
+    } else {
+        buf_a[buf_ib].qs[iqs] = 0;
+        buf_a[buf_ib].qs[iqs + 4] = 0;
+
+        if (iqs == 0) {
+            buf_a[buf_ib].d = FLOAT_TYPE(0.0f);
+        }
     }
 }
 
@@ -218,19 +277,28 @@ ACC_TYPE mmq_dot_product(const uint ib_a) {
 #endif
 
 #if defined(DATA_A_IQ4_XS)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
-    const uint ib_k = ib / 8;
-    const uint ib32 = ib % 8;
-    const uint32_t vui = data_a_packed32[ib_k].qs[4 * ib32 + iqs];
-    const i32vec2 qs = iq4nl_to_i8x8(vui);
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+    if (is_in_bounds) {
+        const uint ib_k = ib / 8;
+        const uint ib32 = ib % 8;
+        const uint32_t vui = data_a_packed32[ib_k].qs[4 * ib32 + iqs];
+        const i32vec2 qs = iq4nl_to_i8x8(vui);
 
-    buf_a[buf_ib].qs[iqs    ] = qs.x;
-    buf_a[buf_ib].qs[iqs + 4] = qs.y;
+        buf_a[buf_ib].qs[iqs    ] = qs.x;
+        buf_a[buf_ib].qs[iqs + 4] = qs.y;
 
-    if (iqs == 0) {
-        const uint sl = (data_a_packed32[ib_k].scales_l >> (4 * ib32)) & 0xF;
-        const uint sh = (data_a_packed32[ib_k].scales_h >> (2 * ib32)) & 3;
-        buf_a[buf_ib].d = FLOAT_TYPE(float(data_a[ib_k].d) * float(int(sl | (sh << 4)) - 32));
+        if (iqs == 0) {
+            const uint sl = (data_a_packed32[ib_k].scales_l >> (4 * ib32)) & 0xF;
+            const uint sh = (data_a_packed32[ib_k].scales_h >> (2 * ib32)) & 3;
+            buf_a[buf_ib].d = FLOAT_TYPE(float(data_a[ib_k].d) * float(int(sl | (sh << 4)) - 32));
+        }
+    } else {
+        buf_a[buf_ib].qs[iqs] = 0;
+        buf_a[buf_ib].qs[iqs + 4] = 0;
+
+        if (iqs == 0) {
+            buf_a[buf_ib].d = FLOAT_TYPE(0.0f);
+        }
     }
 }
 
@@ -256,24 +324,33 @@ ACC_TYPE mmq_dot_product(const uint ib_a) {
 // iqs still refers to a 32-bit integer, meaning 0..7 for 32-wide quants
 #if defined(DATA_A_Q2_K)
 // 4-byte loads for Q2_K blocks (84 bytes)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
-    const uint ib_k = ib / 8;
-    const uint iqs_k = (ib % 8) * 8 + iqs * QUANT_R_MMQ;
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+    if (is_in_bounds) {
+        const uint ib_k = ib / 8;
+        const uint iqs_k = (ib % 8) * 8 + iqs * QUANT_R_MMQ;
 
-    const uint qs_idx = (iqs_k / 32) * 8 + (iqs_k % 8);
-    const uint qs_shift = ((iqs_k % 32) / 8) * 2;
+        const uint qs_idx = (iqs_k / 32) * 8 + (iqs_k % 8);
+        const uint qs_shift = ((iqs_k % 32) / 8) * 2;
 
-    // Repack 4x4 quants into one int
-    const uint32_t vals0 = (data_a_packed32[ib_k].qs[qs_idx    ] >> qs_shift) & 0x03030303;
-    const uint32_t vals1 = (data_a_packed32[ib_k].qs[qs_idx + 1] >> qs_shift) & 0x03030303;
-    const uint32_t vals2 = (data_a_packed32[ib_k].qs[qs_idx + 2] >> qs_shift) & 0x03030303;
-    const uint32_t vals3 = (data_a_packed32[ib_k].qs[qs_idx + 3] >> qs_shift) & 0x03030303;
+        // Repack 4x4 quants into one int
+        const uint32_t vals0 = (data_a_packed32[ib_k].qs[qs_idx    ] >> qs_shift) & 0x03030303;
+        const uint32_t vals1 = (data_a_packed32[ib_k].qs[qs_idx + 1] >> qs_shift) & 0x03030303;
+        const uint32_t vals2 = (data_a_packed32[ib_k].qs[qs_idx + 2] >> qs_shift) & 0x03030303;
+        const uint32_t vals3 = (data_a_packed32[ib_k].qs[qs_idx + 3] >> qs_shift) & 0x03030303;
 
-    buf_a[buf_ib].qs[iqs] = vals0 | (vals1 << 2) | (vals2 << 4) | (vals3 << 6);
+        buf_a[buf_ib].qs[iqs] = vals0 | (vals1 << 2) | (vals2 << 4) | (vals3 << 6);
 
-    if (iqs == 0) {
-        buf_a[buf_ib].dm = FLOAT_TYPEV2(data_a_packed32[ib_k].dm);
-        buf_a[buf_ib].scales = unpack8(uint32_t(data_a_packed16[ib_k].scales[iqs_k / 8])).xy; // vec4 used due to #12147
+        if (iqs == 0) {
+            buf_a[buf_ib].dm = FLOAT_TYPEV2(data_a_packed32[ib_k].dm);
+            buf_a[buf_ib].scales = unpack8(uint32_t(data_a_packed16[ib_k].scales[iqs_k / 8])).xy; // vec4 used due to #12147
+        }
+    } else {
+        buf_a[buf_ib].qs[iqs] = 0;
+
+        if (iqs == 0) {
+            buf_a[buf_ib].dm = FLOAT_TYPEV2(0.0f);
+            buf_a[buf_ib].scales = u8vec2(0);
+        }
     }
 }
 
@@ -305,35 +382,43 @@ ACC_TYPE mmq_dot_product(const uint ib_a) {
 
 #if defined(DATA_A_Q3_K)
 // 2-byte loads for Q3_K blocks (110 bytes)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
-    const uint ib_k = ib / 8;
-    const uint hm_idx = iqs * QUANT_R_MMQ;
-    const uint iqs_k = (ib % 8) * 8 + hm_idx;
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+    if (is_in_bounds) {
+        const uint ib_k = ib / 8;
+        const uint hm_idx = iqs * QUANT_R_MMQ;
+        const uint iqs_k = (ib % 8) * 8 + hm_idx;
 
-    const uint qs_idx = (iqs_k / 32) * 8 + (iqs_k % 8);
-    const uint qs_shift = ((iqs_k % 32) / 8) * 2;
-    const uint hm_shift = iqs_k / 8;
+        const uint qs_idx = (iqs_k / 32) * 8 + (iqs_k % 8);
+        const uint qs_shift = ((iqs_k % 32) / 8) * 2;
+        const uint hm_shift = iqs_k / 8;
 
-    // Repack 2x4 quants into one int
-    // Add the 3rd bit instead of subtracting it to allow packing the quants
-    // vec4 for unpack8 used due to #12147
-    const i8vec2 vals00 = unpack8(int32_t(int16_t((data_a_packed16[ib_k].qs[qs_idx * 2        ] >> qs_shift) & uint16_t(0x0303)))).xy |
-                          unpack8(int32_t(int16_t(((data_a_packed16[ib_k].hmask[hm_idx * 2    ] >> hm_shift) & uint16_t(0x0101))) << 2)).xy;
-    const i8vec2 vals01 = unpack8(int32_t(int16_t((data_a_packed16[ib_k].qs[qs_idx * 2 + 1    ] >> qs_shift) & uint16_t(0x0303)))).xy |
-                          unpack8(int32_t(int16_t(((data_a_packed16[ib_k].hmask[hm_idx * 2 + 1] >> hm_shift) & uint16_t(0x0101))) << 2)).xy;
-    const i8vec2 vals10 = unpack8(int32_t(int16_t((data_a_packed16[ib_k].qs[qs_idx * 2 + 2    ] >> qs_shift) & uint16_t(0x0303)))).xy |
-                          unpack8(int32_t(int16_t(((data_a_packed16[ib_k].hmask[hm_idx * 2 + 2] >> hm_shift) & uint16_t(0x0101))) << 2)).xy;
-    const i8vec2 vals11 = unpack8(int32_t(int16_t((data_a_packed16[ib_k].qs[qs_idx * 2 + 3    ] >> qs_shift) & uint16_t(0x0303)))).xy |
-                          unpack8(int32_t(int16_t(((data_a_packed16[ib_k].hmask[hm_idx * 2 + 3] >> hm_shift) & uint16_t(0x0101))) << 2)).xy;
-    buf_a[buf_ib].qs[iqs] = pack32(u8vec4(vals00.x, vals00.y, vals01.x, vals01.y)) |
-                           (pack32(u8vec4(vals10.x, vals10.y, vals11.x, vals11.y)) << 4);
+        // Repack 2x4 quants into one int
+        // Add the 3rd bit instead of subtracting it to allow packing the quants
+        // vec4 for unpack8 used due to #12147
+        const i8vec2 vals00 = unpack8(int32_t(int16_t((data_a_packed16[ib_k].qs[qs_idx * 2        ] >> qs_shift) & uint16_t(0x0303)))).xy |
+                              unpack8(int32_t(int16_t(((data_a_packed16[ib_k].hmask[hm_idx * 2    ] >> hm_shift) & uint16_t(0x0101))) << 2)).xy;
+        const i8vec2 vals01 = unpack8(int32_t(int16_t((data_a_packed16[ib_k].qs[qs_idx * 2 + 1    ] >> qs_shift) & uint16_t(0x0303)))).xy |
+                              unpack8(int32_t(int16_t(((data_a_packed16[ib_k].hmask[hm_idx * 2 + 1] >> hm_shift) & uint16_t(0x0101))) << 2)).xy;
+        const i8vec2 vals10 = unpack8(int32_t(int16_t((data_a_packed16[ib_k].qs[qs_idx * 2 + 2    ] >> qs_shift) & uint16_t(0x0303)))).xy |
+                              unpack8(int32_t(int16_t(((data_a_packed16[ib_k].hmask[hm_idx * 2 + 2] >> hm_shift) & uint16_t(0x0101))) << 2)).xy;
+        const i8vec2 vals11 = unpack8(int32_t(int16_t((data_a_packed16[ib_k].qs[qs_idx * 2 + 3    ] >> qs_shift) & uint16_t(0x0303)))).xy |
+                              unpack8(int32_t(int16_t(((data_a_packed16[ib_k].hmask[hm_idx * 2 + 3] >> hm_shift) & uint16_t(0x0101))) << 2)).xy;
+        buf_a[buf_ib].qs[iqs] = pack32(u8vec4(vals00.x, vals00.y, vals01.x, vals01.y)) |
+                               (pack32(u8vec4(vals10.x, vals10.y, vals11.x, vals11.y)) << 4);
 
-    if (iqs == 0) {
-        const uint is = iqs_k / 4;
-        const i8vec2 scales = i8vec2(unpack8(uint32_t(((data_a_packed16[ib_k].scales[(is % 8      ) / 2] >> (4 * (is / 8))) & 0x0F0F) |
-                                                     (((data_a_packed16[ib_k].scales[(8 + (is % 4)) / 2] >> (2 * (is / 4))) & 0x0303) << 4))).xy); // vec4 used due to #12147
+        if (iqs == 0) {
+            const uint is = iqs_k / 4;
+            const i8vec2 scales = i8vec2(unpack8(uint32_t(((data_a_packed16[ib_k].scales[(is % 8      ) / 2] >> (4 * (is / 8))) & 0x0F0F) |
+                                                         (((data_a_packed16[ib_k].scales[(8 + (is % 4)) / 2] >> (2 * (is / 4))) & 0x0303) << 4))).xy); // vec4 used due to #12147
 
-        buf_a[buf_ib].d_scales = FLOAT_TYPEV2(float(data_a_packed16[ib_k].d) * vec2(scales - 32));
+            buf_a[buf_ib].d_scales = FLOAT_TYPEV2(float(data_a_packed16[ib_k].d) * vec2(scales - 32));
+        }
+    } else {
+        buf_a[buf_ib].qs[iqs] = 0;
+
+        if (iqs == 0) {
+            buf_a[buf_ib].d_scales = FLOAT_TYPEV2(0.0f);
+        }
     }
 }
 
@@ -371,39 +456,47 @@ ACC_TYPE mmq_dot_product(const uint ib_a) {
 
 #if defined(DATA_A_Q4_K) || defined(DATA_A_Q5_K)
 // 4-byte loads for Q4_K blocks (144 bytes) and Q5_K blocks (176 bytes)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
-    const uint ib_k = ib / 8;
-    const uint iqs_k = (ib % 8) * 8 + iqs * QUANT_R_MMQ;
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+    if (is_in_bounds) {
+        const uint ib_k = ib / 8;
+        const uint iqs_k = (ib % 8) * 8 + iqs * QUANT_R_MMQ;
 
-    const uint qs_idx = (iqs_k / 16) * 8 + (iqs_k % 8);
-    const uint qs_shift = ((iqs_k % 16) / 8) * 4;
+        const uint qs_idx = (iqs_k / 16) * 8 + (iqs_k % 8);
+        const uint qs_shift = ((iqs_k % 16) / 8) * 4;
 
-    // Repack 2x4 quants into one int
-#if defined(DATA_A_Q4_K)
-    const uint32_t vals0 = (data_a_packed32[ib_k].qs[qs_idx    ] >> qs_shift) & 0x0F0F0F0F;
-    const uint32_t vals1 = (data_a_packed32[ib_k].qs[qs_idx + 1] >> qs_shift) & 0x0F0F0F0F;
+        // Repack 2x4 quants into one int
+    #if defined(DATA_A_Q4_K)
+        const uint32_t vals0 = (data_a_packed32[ib_k].qs[qs_idx    ] >> qs_shift) & 0x0F0F0F0F;
+        const uint32_t vals1 = (data_a_packed32[ib_k].qs[qs_idx + 1] >> qs_shift) & 0x0F0F0F0F;
 
-    buf_a[buf_ib].qs[iqs] = vals0 | (vals1 << 4);
-#else // defined(DATA_A_Q5_K)
-    const uint qh_idx = iqs * QUANT_R_MMQ;
-    const uint qh_shift = iqs_k / 8;
+        buf_a[buf_ib].qs[iqs] = vals0 | (vals1 << 4);
+    #else // defined(DATA_A_Q5_K)
+        const uint qh_idx = iqs * QUANT_R_MMQ;
+        const uint qh_shift = iqs_k / 8;
 
-    buf_a[buf_ib].qs[iqs] = int32_t(((data_a_packed32[ib_k].qs[qs_idx] >> qs_shift) & 0x0F0F0F0F) |
-                                   (((data_a_packed32[ib_k].qh[qh_idx] >> qh_shift) & 0x01010101) << 4));
-#endif
+        buf_a[buf_ib].qs[iqs] = int32_t(((data_a_packed32[ib_k].qs[qs_idx] >> qs_shift) & 0x0F0F0F0F) |
+                                       (((data_a_packed32[ib_k].qh[qh_idx] >> qh_shift) & 0x01010101) << 4));
+    #endif
 
-    if (iqs == 0) {
-        // Scale index
-        const uint is = iqs_k / 8;
-        u8vec2 scale_dm;
-        if (is < 4) {
-            scale_dm = u8vec2(data_a[ib_k].scales[is] & 0x3F, data_a[ib_k].scales[is + 4] & 0x3F);
-        } else {
-            scale_dm = u8vec2((data_a[ib_k].scales[is+4] & 0xF) | ((data_a[ib_k].scales[is-4] & 0xC0) >> 2),
-                              (data_a[ib_k].scales[is+4] >>  4) | ((data_a[ib_k].scales[is  ] & 0xC0) >> 2));
+        if (iqs == 0) {
+            // Scale index
+            const uint is = iqs_k / 8;
+            u8vec2 scale_dm;
+            if (is < 4) {
+                scale_dm = u8vec2(data_a[ib_k].scales[is] & 0x3F, data_a[ib_k].scales[is + 4] & 0x3F);
+            } else {
+                scale_dm = u8vec2((data_a[ib_k].scales[is+4] & 0xF) | ((data_a[ib_k].scales[is-4] & 0xC0) >> 2),
+                                  (data_a[ib_k].scales[is+4] >>  4) | ((data_a[ib_k].scales[is  ] & 0xC0) >> 2));
+            }
+
+            buf_a[buf_ib].dm = FLOAT_TYPEV2(vec2(data_a_packed32[ib_k].dm) * vec2(scale_dm));
         }
+    } else {
+        buf_a[buf_ib].qs[iqs] = 0;
 
-        buf_a[buf_ib].dm = FLOAT_TYPEV2(vec2(data_a_packed32[ib_k].dm) * vec2(scale_dm));
+        if (iqs == 0) {
+            buf_a[buf_ib].dm = FLOAT_TYPEV2(0.0f);
+        }
     }
 }
 
@@ -434,27 +527,35 @@ ACC_TYPE mmq_dot_product(const uint ib_a) {
 
 #if defined(DATA_A_Q6_K)
 // 2-byte loads for Q6_K blocks (210 bytes)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
-    const uint ib_k = ib / 8;
-    const uint iqs_k = (ib % 8) * 8 + iqs;
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+    if (is_in_bounds) {
+        const uint ib_k = ib / 8;
+        const uint iqs_k = (ib % 8) * 8 + iqs;
 
-    const uint ql_idx = (iqs_k / 32) * 16 + iqs_k % 16;
-    const uint ql_shift = ((iqs_k % 32) / 16) * 4;
+        const uint ql_idx = (iqs_k / 32) * 16 + iqs_k % 16;
+        const uint ql_shift = ((iqs_k % 32) / 16) * 4;
 
-    const uint qh_idx = (iqs_k / 32) * 8 + iqs;
-    const uint qh_shift = ((iqs_k % 32) / 8) * 2;
+        const uint qh_idx = (iqs_k / 32) * 8 + iqs;
+        const uint qh_shift = ((iqs_k % 32) / 8) * 2;
 
-    const i8vec2 vals00 = (unpack8(int32_t((data_a_packed16[ib_k].ql[ql_idx * 2    ] >> ql_shift) & uint16_t(0x0F0F))).xy |
-                          unpack8(int32_t(((data_a_packed16[ib_k].qh[qh_idx * 2    ] >> qh_shift) & uint16_t(0x0303)) << 4)).xy) - int8_t(32);
-    const i8vec2 vals01 = (unpack8(int32_t((data_a_packed16[ib_k].ql[ql_idx * 2 + 1] >> ql_shift) & uint16_t(0x0F0F))).xy |
-                          unpack8(int32_t(((data_a_packed16[ib_k].qh[qh_idx * 2 + 1] >> qh_shift) & uint16_t(0x0303)) << 4)).xy) - int8_t(32);
-    buf_a[buf_ib].qs[iqs] = pack32(i8vec4(vals00.x, vals00.y, vals01.x, vals01.y));
+        const i8vec2 vals00 = (unpack8(int32_t((data_a_packed16[ib_k].ql[ql_idx * 2    ] >> ql_shift) & uint16_t(0x0F0F))).xy |
+                              unpack8(int32_t(((data_a_packed16[ib_k].qh[qh_idx * 2    ] >> qh_shift) & uint16_t(0x0303)) << 4)).xy) - int8_t(32);
+        const i8vec2 vals01 = (unpack8(int32_t((data_a_packed16[ib_k].ql[ql_idx * 2 + 1] >> ql_shift) & uint16_t(0x0F0F))).xy |
+                              unpack8(int32_t(((data_a_packed16[ib_k].qh[qh_idx * 2 + 1] >> qh_shift) & uint16_t(0x0303)) << 4)).xy) - int8_t(32);
+        buf_a[buf_ib].qs[iqs] = pack32(i8vec4(vals00.x, vals00.y, vals01.x, vals01.y));
 
-    if (iqs == 0) {
-        const uint is = iqs_k / 4;
-        const i8vec2 scales = unpack8(int32_t(data_a_packed16[ib_k].scales[is / 2])).xy;
+        if (iqs == 0) {
+            const uint is = iqs_k / 4;
+            const i8vec2 scales = unpack8(int32_t(data_a_packed16[ib_k].scales[is / 2])).xy;
 
-        buf_a[buf_ib].d_scales = FLOAT_TYPEV2(float(data_a_packed16[ib_k].d) * vec2(scales));
+            buf_a[buf_ib].d_scales = FLOAT_TYPEV2(float(data_a_packed16[ib_k].d) * vec2(scales));
+        }
+    } else {
+        buf_a[buf_ib].qs[iqs] = 0;
+
+        if (iqs == 0) {
+            buf_a[buf_ib].d_scales = FLOAT_TYPEV2(0.0f);
+        }
     }
 }
 
@@ -491,32 +592,41 @@ ACC_TYPE mmq_dot_product(const uint ib_a) {
 
 #if defined(DATA_A_IQ3_S)
 // 2-byte loads for IQ3_S blocks (110 bytes)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
-    const uint ib_k = ib / 8;
-    const uint ib32 = ib % 8;
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+    if (is_in_bounds) {
+        const uint ib_k = ib / 8;
+        const uint ib32 = ib % 8;
 
-    // grid indices for qs[2 * iqs] and qs[2 * iqs + 1]
-    const uint qs    = uint(data_a_packed16[ib_k].qs[ib32 * 4 + iqs]);
-    // their two high index bits
-    const uint qh    = uint(data_a_packed16[ib_k].qh[ib32 / 2]) >> ((ib32 & 1) * 8 + 2 * iqs);
-    // one sign bit per value, 8 values
-    const uint signs = uint(data_a_packed16[ib_k].signs[ib32 * 2 + iqs / 2]) >> ((iqs & 1) * 8);
+        // grid indices for qs[2 * iqs] and qs[2 * iqs + 1]
+        const uint qs    = uint(data_a_packed16[ib_k].qs[ib32 * 4 + iqs]);
+        // their two high index bits
+        const uint qh    = uint(data_a_packed16[ib_k].qh[ib32 / 2]) >> ((ib32 & 1) * 8 + 2 * iqs);
+        // one sign bit per value, 8 values
+        const uint signs = uint(data_a_packed16[ib_k].signs[ib32 * 2 + iqs / 2]) >> ((iqs & 1) * 8);
 
-    // grid holds 4 values of 1..15, one per byte
-    const ivec4 vals0 = ivec4(unpack8(iq3s_grid[( qs       & 0xFF) | ((qh & 1) << 8)]));
-    const ivec4 vals1 = ivec4(unpack8(iq3s_grid[((qs >> 8) & 0xFF) | ((qh & 2) << 7)]));
+        // grid holds 4 values of 1..15, one per byte
+        const ivec4 vals0 = ivec4(unpack8(iq3s_grid[( qs       & 0xFF) | ((qh & 1) << 8)]));
+        const ivec4 vals1 = ivec4(unpack8(iq3s_grid[((qs >> 8) & 0xFF) | ((qh & 2) << 7)]));
 
-    // negate with (v ^ -s) - -s to avoid branches
-    const ivec4 m0 = -(ivec4(signs,      signs >> 1, signs >> 2, signs >> 3) & 1);
-    const ivec4 m1 = -(ivec4(signs >> 4, signs >> 5, signs >> 6, signs >> 7) & 1);
+        // negate with (v ^ -s) - -s to avoid branches
+        const ivec4 m0 = -(ivec4(signs,      signs >> 1, signs >> 2, signs >> 3) & 1);
+        const ivec4 m1 = -(ivec4(signs >> 4, signs >> 5, signs >> 6, signs >> 7) & 1);
 
-    buf_a[buf_ib].qs[2 * iqs    ] = pack32(i8vec4((vals0 ^ m0) - m0));
-    buf_a[buf_ib].qs[2 * iqs + 1] = pack32(i8vec4((vals1 ^ m1) - m1));
+        buf_a[buf_ib].qs[2 * iqs    ] = pack32(i8vec4((vals0 ^ m0) - m0));
+        buf_a[buf_ib].qs[2 * iqs + 1] = pack32(i8vec4((vals1 ^ m1) - m1));
 
-    if (iqs == 0) {
-        const uint scale = (uint(data_a_packed16[ib_k].scales[ib32 / 4]) >> ((ib32 & 3) * 4)) & 0xF;
+        if (iqs == 0) {
+            const uint scale = (uint(data_a_packed16[ib_k].scales[ib32 / 4]) >> ((ib32 & 3) * 4)) & 0xF;
 
-        buf_a[buf_ib].d = FLOAT_TYPE(float(data_a_packed16[ib_k].d) * float(1 + 2 * scale));
+            buf_a[buf_ib].d = FLOAT_TYPE(float(data_a_packed16[ib_k].d) * float(1 + 2 * scale));
+        }
+    } else {
+        buf_a[buf_ib].qs[2 * iqs] = 0;
+        buf_a[buf_ib].qs[2 * iqs + 1] = 0;
+
+        if (iqs == 0) {
+            buf_a[buf_ib].d = FLOAT_TYPE(0.0f);
+        }
     }
 }
 

ggml/src/ggml-zdnn/common.hpp

@@ -6,11 +6,17 @@
 
 #include "zdnn.h"
 
-#include <vector>
+#include <algorithm>
+#include <cstdlib>
+#include <future>
+#include <map>
 #include <memory>
+#include <utility>
+#include <vector>
 
-#define GGML_ZDNN_NAME    "zDNN"
-#define GGML_ZDNN_VERSION ZDNN_VERNUM
+#define GGML_ZDNN_NAME          "zDNN"
+#define GGML_ZDNN_VERSION       ZDNN_VERNUM
+#define GGML_ZDNN_SCRATCH_SLOTS 3  // int8 inputs, DFloat16 partial results, partial results for grouped matmul
 
 #define ZDNN_CHECK(stmt)                \
     do {                                \
@@ -30,11 +36,6 @@ struct ggml_backend_zdnn_device_context {
     char name[128];
 };
 
-struct ggml_backend_zdnn_context {
-    int device;
-    ggml_cgraph * gf;
-};
-
 struct ggml_backend_zdnn_buffer {
     void * data;
     ggml_backend_zdnn_buffer * extra;  // for bias, etc.
@@ -44,6 +45,9 @@ struct ggml_backend_zdnn_buffer {
     zdnn_tensor_desc tfm_desc;
     zdnn_ztensor     ztensor;
 
+    std::vector<float> scales;          // per-row scales for quantized weights
+    std::vector<float> channel_scales;  // per-input-channel scales for quantized weights, applied to the inputs
+
     char name[GGML_MAX_NAME];
 };
 
@@ -56,4 +60,75 @@ struct ggml_backend_zdnn_buffer_context {
     std::vector<std::unique_ptr<ggml_backend_zdnn_buffer>> buffers;
 };
 
+struct ggml_backend_zdnn_context {
+    int device;
+    int32_t n_threads;  // for the CPU work around zDNN calls
+    ggml_cgraph * gf;
+
+    // zero bias for stacked quantized matmuls, keyed by (n_stacks, ne0)
+    std::map<std::pair<int64_t, int64_t>, std::unique_ptr<ggml_backend_zdnn_buffer>> zero_bias;
+
+    std::vector<int8_t> scratch_q;
+    std::vector<float>  scratch_f32;
+
+    void * scratch_ztensor[GGML_ZDNN_SCRATCH_SLOTS]      = {};  // 4K aligned
+    size_t scratch_ztensor_size[GGML_ZDNN_SCRATCH_SLOTS] = {};
+};
+
+template <typename T>
+static inline T * ggml_zdnn_resize_scratch_max(std::vector<T> & scratch, int64_t n) {
+    if ((int64_t)scratch.size() < n) {
+        scratch.resize(n);
+    }
+
+    return scratch.data();
+}
+
+//
+// Scratch buffers
+//
+
+static inline void ggml_zdnn_init_scratch_ztensor(ggml_backend_zdnn_context * ctx,
+                                                  zdnn_tensor_desc          * pre_tfm_desc,
+                                                  zdnn_tensor_desc          * tfm_desc,
+                                                  zdnn_ztensor              * ztensor,
+                                                  int32_t                     slot) {
+
+    const uint64_t size = zdnn_getsize_ztensor(tfm_desc);
+
+    if (ctx->scratch_ztensor_size[slot] < size) {
+        free(ctx->scratch_ztensor[slot]);
+        ctx->scratch_ztensor[slot]      = std::aligned_alloc(4096, GGML_PAD(size, 4096));
+        ctx->scratch_ztensor_size[slot] = GGML_PAD(size, 4096);
+        GGML_ASSERT(ctx->scratch_ztensor[slot] != nullptr);
+    }
+
+    zdnn_init_quantized_ztensor(pre_tfm_desc, tfm_desc, 1.0f, 0.0f, ztensor);
+    ztensor->buffer      = ctx->scratch_ztensor[slot];
+    ztensor->buffer_size = size;
+}
+
+//
+// Threading
+//
+
+// split [0, n) into n_threads ranges, the calling thread runs the first range
+template <typename F>
+static inline void ggml_zdnn_parallel_for(int32_t n_threads, int64_t n, const F & f) {
+    n_threads = (int32_t)std::max<int64_t>(1, std::min<int64_t>(n_threads, n));
+
+    std::vector<std::future<void>> tasks;
+    for (int32_t ith = 1; ith < n_threads; ith++) {
+        const int64_t start = (ith + 0) * n / n_threads;
+        const int64_t end   = (ith + 1) * n / n_threads;
+        tasks.push_back(std::async(std::launch::async, [&f, start, end]() { f(start, end); }));
+    }
+
+    f(0, n/n_threads);
+
+    for (auto & task : tasks) {
+        task.get();
+    }
+}
+
 #endif  // GGML_ZDNN_COMMON_HPP

ggml/src/ggml-zdnn/ggml-zdnn.cpp

@@ -4,6 +4,7 @@
 
 #include "ggml-zdnn/common.hpp"
 #include "ggml-zdnn/mmf.hpp"
+#include "ggml-zdnn/mmq.hpp"
 #include "ggml-zdnn/utils.hpp"
 #include "ggml.h"
 
@@ -13,15 +14,17 @@
 #include <unistd.h>
 
 static void ggml_zdnn_compute_forward_mul_mat(
-    const ggml_backend_zdnn_context * ctx,
+    ggml_backend_zdnn_context * ctx,
           ggml_tensor * dst) {
 
     const ggml_tensor * src0 = dst->src[0];  // weights
     const ggml_tensor * src1 = dst->src[1];  // inputs
 
-    // TODO: implement support for quantized types
-    // we currently only support f32, f16, and bf16
-    ggml_zdnn_mul_mat_f(ctx, src0, src1, dst);
+    if (ggml_is_quantized(src0->type)) {
+        ggml_zdnn_mul_mat_q(ctx, src0, src1, dst);
+    } else {
+        ggml_zdnn_mul_mat_f(ctx, src0, src1, dst);
+    }
 }
 
 static bool ggml_zdnn_compute_forward(
@@ -98,7 +101,7 @@ static bool ggml_zdnn_supports_op(const ggml_backend_zdnn_device_context * ctx_d
 
                 if (!ggml_is_matrix(weights) || !ggml_is_matrix(inputs) ||
                     !ggml_is_contiguous(weights) || !ggml_is_contiguous(inputs) ||
-                    weights->view_src != nullptr || inputs->view_src != nullptr ||
+                    weights->view_src != nullptr ||
                     ne0 > max_batch || ne1 > max_batch || ne10 > max_batch) {
                         return false;
                 }
@@ -107,7 +110,11 @@ static bool ggml_zdnn_supports_op(const ggml_backend_zdnn_device_context * ctx_d
                     case GGML_TYPE_F32:
                     case GGML_TYPE_F16:
                     case GGML_TYPE_BF16:
-                        return true;
+                        // the inputs ztensor is used, and views have none
+                        return inputs->view_src == nullptr;
+                    case GGML_TYPE_Q8_0:
+                        // the inputs are quantized from their data, so views work as well
+                        return ctx_dev->has_parmblkformat_1 && inputs->type == GGML_TYPE_F32;
                     default:
                         return false;
                 }
@@ -186,6 +193,7 @@ static ggml_backend_zdnn_context * ggml_zdnn_init(ggml_backend_dev_t dev) {
     GGML_LOG_INFO("%s: NNPA_PARMBLKFORMAT_0 = %s\n", __func__, ctx_dev->has_parmblkformat_0 ? "true" : "false");
     GGML_LOG_INFO("%s: NNPA_PARMBLKFORMAT_1 = %s\n", __func__, ctx_dev->has_parmblkformat_1 ? "true" : "false");
 
+    ctx->n_threads = GGML_DEFAULT_N_THREADS;
     ctx->gf = nullptr;
 
     return ctx;
@@ -193,6 +201,14 @@ static ggml_backend_zdnn_context * ggml_zdnn_init(ggml_backend_dev_t dev) {
 
 static void ggml_zdnn_free(ggml_backend_zdnn_context * ctx) {
     GGML_LOG_INFO("%s: deallocating\n", __func__);
+
+    for (const auto & it : ctx->zero_bias) {
+        ZDNN_CHECK(zdnn_free_ztensor_buffer(&it.second->ztensor));
+    }
+    for (void * scratch : ctx->scratch_ztensor) {
+        free(scratch);
+    }
+
     delete ctx;
 }
 
@@ -295,9 +311,19 @@ static void ggml_backend_zdnn_buffer_set_tensor(ggml_backend_buffer_t buffer, gg
 
     ggml_backend_zdnn_buffer * extra = (ggml_backend_zdnn_buffer *)tensor->extra;
 
+    if (ggml_is_quantized(tensor->type)) {
+        GGML_ASSERT(offset == 0 && size == ggml_nbytes(tensor));
+        ggml_zdnn_load_tensor_q(extra, tensor);
+        return;
+    }
+
     // Fixes the LLAMA_SET_ROWS bug
     // see: https://github.com/ggml-org/llama.cpp/issues/15414
-    if (tensor->buffer->usage == GGML_BACKEND_BUFFER_USAGE_COMPUTE && extra->ztensor.is_transformed) zdnn_reset_ztensor(&extra->ztensor);
+    // compute tensors are transformed on first use instead, as not every op reads the ztensor
+    if (tensor->buffer->usage == GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
+        zdnn_reset_ztensor(&extra->ztensor);
+        return;
+    }
     if (extra->ztensor.is_transformed == false) ggml_zdnn_load_tensor(extra->ztensor, tensor->data);
 
     GGML_UNUSED(buffer);
@@ -621,10 +647,20 @@ static ggml_backend_feature * ggml_backend_zdnn_get_features(ggml_backend_reg_t
     GGML_UNUSED(reg);
 }
 
+static void ggml_backend_zdnn_set_n_threads(ggml_backend_t backend, int n_threads) {
+    GGML_ASSERT(ggml_backend_is_zdnn(backend));
+
+    ggml_backend_zdnn_context * ctx = (ggml_backend_zdnn_context *)backend->context;
+    ctx->n_threads = n_threads;
+}
+
 static void * ggml_backend_zdnn_get_proc_address(ggml_backend_reg_t reg, const char * name) {
     if (strcmp(name, "ggml_backend_get_features") == 0) {
         return (void *) ggml_backend_zdnn_get_features;
     }
+    if (strcmp(name, "ggml_backend_set_n_threads") == 0) {
+        return (void *) ggml_backend_zdnn_set_n_threads;
+    }
 
     return NULL;
 

ggml/src/ggml-zdnn/mmf.cpp

@@ -1,5 +1,6 @@
 #include "ggml.h"
 #include "mmf.hpp"
+#include "utils.hpp"
 
 void ggml_zdnn_mul_mat_f(
     const ggml_backend_zdnn_context * ctx,
@@ -65,6 +66,14 @@ void ggml_zdnn_mul_mat_f(
     GGML_ASSERT(inputs_extra->pre_tfm_desc.dim1  == inputs->ne[0]  && "inputs_extra->pre_tfm_desc.dim1 must match inputs->ne[0]");
     GGML_ASSERT(inputs_extra->pre_tfm_desc.dim2  == inputs->ne[1]  && "inputs_extra->pre_tfm_desc.dim2 must match inputs->ne[1]");
 
+    // compute tensors are transformed on first use
+    if (!weights_extra->ztensor.is_transformed) {
+        ggml_zdnn_load_tensor(weights_extra->ztensor, weights->data);
+    }
+    if (!inputs_extra->ztensor.is_transformed) {
+        ggml_zdnn_load_tensor(inputs_extra->ztensor, inputs->data);
+    }
+
     ZDNN_CHECK(zdnn_matmul_transpose_op(&inputs_extra->ztensor, &weights_extra->ztensor, &bias_extra->ztensor,
                                         false, true, MATMUL_OP_ADDITION, &output_extra->ztensor));
     // TODO: Remove in the future as we are currently DLF16 -> FP32 then in the next op, FP32 -> DLF16 again. Inefficient.

ggml/src/ggml-zdnn/mmq.cpp

@@ -0,0 +1,331 @@
+#include "ggml.h"
+#include "mmq.hpp"
+#include "quantize.hpp"
+
+#include <algorithm>
+#include <future>
+#include <vector>
+
+// view of the stacks [s0, s0 + n) of a stacked ztensor, the stacks are the outermost dim of its buffer
+// the view descs must outlive the view
+static zdnn_ztensor ggml_zdnn_stack_view(zdnn_tensor_desc * pre_tfm_desc,
+                                         zdnn_tensor_desc * tfm_desc,
+                                   const zdnn_ztensor     & src,
+                                         int64_t            n_stacks,
+                                         int64_t            s0,
+                                         int64_t            n) {
+
+    GGML_ASSERT(src.buffer_size % n_stacks == 0);
+    const uint64_t stack_bytes = src.buffer_size / n_stacks;
+
+    *pre_tfm_desc = *src.pre_transformed_desc;
+    *tfm_desc     = *src.transformed_desc;
+
+    switch (pre_tfm_desc->layout) {
+        case ZDNN_3DS:
+            {
+                pre_tfm_desc->dim3 = n;
+            } break;
+        case ZDNN_2DS:
+            {
+                pre_tfm_desc->dim2 = n;
+            } break;
+        default:
+            GGML_ABORT("%s: unsupported stacked layout", __func__);
+    }
+
+    tfm_desc->dim4 = n;
+    GGML_ASSERT(zdnn_getsize_ztensor(tfm_desc) == n*stack_bytes);
+
+    zdnn_ztensor view = src;
+    view.pre_transformed_desc = pre_tfm_desc;
+    view.transformed_desc     = tfm_desc;
+    view.buffer               = (char *) src.buffer + s0*stack_bytes;
+    view.buffer_size          = n * stack_bytes;
+    return view;
+}
+
+// zero bias of n_stacks stacks, created once per shape as it never changes
+static const zdnn_ztensor & ggml_zdnn_zero_bias(ggml_backend_zdnn_context * ctx, int64_t n_stacks, int64_t ne0) {
+    std::unique_ptr<ggml_backend_zdnn_buffer> & bias = ctx->zero_bias[{ n_stacks, ne0 }];
+
+    if (!bias) {
+        bias = std::make_unique<ggml_backend_zdnn_buffer>();
+
+        std::vector<float> zeros(n_stacks * ne0, 0.0f);
+
+        zdnn_init_pre_transformed_desc(ZDNN_2DS, FP32, &bias->pre_tfm_desc, n_stacks, ne0);
+        ZDNN_CHECK(zdnn_generate_transformed_desc(&bias->pre_tfm_desc, &bias->tfm_desc));
+        ZDNN_CHECK(zdnn_init_ztensor_with_malloc(&bias->pre_tfm_desc, &bias->tfm_desc, &bias->ztensor));
+        ZDNN_CHECK(zdnn_transform_ztensor(&bias->ztensor, zeros.data()));
+    }
+
+    return bias->ztensor;
+}
+
+// the weights are stacked into groups along ne00, and each group of each row of the inputs and weights has its own scale
+// zDNN multiplies the groups in stacked calls, then the partial results of each group are scaled and summed here
+static void ggml_zdnn_mul_mat_q_grouped(ggml_backend_zdnn_context * ctx,
+                                        int32_t                     n_threads,
+                                  const ggml_tensor               * inputs,
+                                  const ggml_backend_zdnn_buffer  * weights_extra,
+                                        ggml_tensor               * output) {
+    const int64_t ne10 = inputs->ne[0];
+    const int64_t ne11 = inputs->ne[1];
+    const int64_t ne0  = output->ne[0];
+
+    const int64_t n_groups = weights_extra->pre_tfm_desc.dim3;
+    const int64_t gs       = weights_extra->pre_tfm_desc.dim2;
+
+    // the last group is zero padded when gs does not divide ne10
+    GGML_ASSERT(n_groups * gs >= ne10 && (n_groups - 1) * gs < ne10);
+
+    // the rounding errors of the inputs are stacked after the inputs, see ggml_zdnn_quantize_inputs
+    int8_t * inputs_q = ggml_zdnn_resize_scratch_max(ctx->scratch_q, 2 * n_groups * gs * ne11);
+    std::vector<float> inputs_scales(n_groups * ne11);
+    ggml_zdnn_quantize_inputs(n_threads, inputs, gs, weights_extra->channel_scales.data(),
+                            inputs_q, inputs_q + n_groups * gs * ne11, inputs_scales.data());
+
+    zdnn_tensor_desc inputs_pre_tfm_desc;
+    zdnn_tensor_desc inputs_tfm_desc;
+    zdnn_ztensor     inputs_ztensor;
+
+    zdnn_init_pre_transformed_desc(ZDNN_3DS, INT8, &inputs_pre_tfm_desc, 2 * n_groups, ne11, gs);
+    ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&inputs_pre_tfm_desc, QUANTIZED_INT8, &inputs_tfm_desc));
+    ggml_zdnn_init_scratch_ztensor(ctx, &inputs_pre_tfm_desc, &inputs_tfm_desc, &inputs_ztensor, 0);
+    ZDNN_CHECK(zdnn_transform_quantized_ztensor(&inputs_ztensor, false, -127, 127, inputs_q));
+
+    // each group of the partial results is read back as its own 2D ztensor
+    zdnn_tensor_desc group_pre_tfm_desc;
+    zdnn_tensor_desc group_tfm_desc;
+
+    zdnn_init_pre_transformed_desc(ZDNN_2D, FP32, &group_pre_tfm_desc, ne11, ne0);
+    ZDNN_CHECK(zdnn_generate_transformed_desc(&group_pre_tfm_desc, &group_tfm_desc));
+
+    const uint64_t group_bytes = zdnn_getsize_ztensor(&group_tfm_desc);
+
+    // stacked calls with a large output are much slower per group, so run the groups in chunks
+    // 8 MiB per chunk was benchmarked to be the best on Llama-3.2-1B Q8_0
+    // always re-measure with different models to see if different sizes need a dispatching logic
+    constexpr int64_t chunk_bytes = 8*1024*1024;
+
+    const int64_t n_chunk = std::clamp<int64_t>(chunk_bytes / (int64_t)group_bytes, 1, n_groups);
+
+    // all zeros, so it is already the pre-computed bias for every chunk
+    const zdnn_ztensor & bias_ztensor = ggml_zdnn_zero_bias(ctx, n_chunk, ne0);
+
+    // two partial buffers, so the next chunk runs on the NNPA while this one is summed
+    // each holds the results of a chunk of inputs, followed by the results of their rounding errors
+    zdnn_tensor_desc partial_pre_tfm_desc;
+    zdnn_tensor_desc partial_tfm_desc;
+    zdnn_ztensor     partial_ztensor[2];
+
+    zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &partial_pre_tfm_desc, 2 * n_chunk, ne11, ne0);
+    ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&partial_pre_tfm_desc, QUANTIZED_DLFLOAT16, &partial_tfm_desc));
+    GGML_ASSERT(zdnn_getsize_ztensor(&partial_tfm_desc) == (uint64_t)(2 * n_chunk) * group_bytes);
+    for (int32_t i = 0; i < 2; i++) {
+        ggml_zdnn_init_scratch_ztensor(ctx, &partial_pre_tfm_desc, &partial_tfm_desc, &partial_ztensor[i], 1 + i);
+    }
+
+    const float * weights_scales = weights_extra->scales.data();
+    float * partial     = ggml_zdnn_resize_scratch_max(ctx->scratch_f32, 2 * ne11 * ne0);
+    float * partial_res = partial + ne11 * ne0;
+
+    const auto sum_chunk = [&](const zdnn_ztensor & chunk, int64_t g0, int64_t n) {
+        zdnn_ztensor group_ztensor;
+        zdnn_init_ztensor(&group_pre_tfm_desc, &group_tfm_desc, &group_ztensor);
+
+        for (int64_t g = g0; g < g0 + n; g++) {
+            group_ztensor.buffer_size    = group_bytes;
+            group_ztensor.is_transformed = true;
+
+            group_ztensor.buffer = (char *)chunk.buffer + (g - g0)*group_bytes;
+            ZDNN_CHECK(zdnn_transform_origtensor(&group_ztensor, partial));
+
+            group_ztensor.buffer = (char *)chunk.buffer + (n_chunk + g - g0)*group_bytes;
+            ZDNN_CHECK(zdnn_transform_origtensor(&group_ztensor, partial_res));
+
+            const float * sw = weights_scales + g*ne0;
+            for (int64_t i1 = 0; i1 < ne11; i1++) {
+                const float   sa = inputs_scales[g*ne11 + i1];
+                const float * p  = partial     + i1*ne0;
+                const float * pr = partial_res + i1*ne0;
+                float       * y  = (float *)output->data + i1*ne0;
+
+                if (g == 0) {
+                    for (int64_t i0 = 0; i0 < ne0; i0++) {
+                        y[i0] = sa * sw[i0] * (p[i0] + pr[i0] * (1.0f/254.0f));
+                    }
+                } else {
+                    for (int64_t i0 = 0; i0 < ne0; i0++) {
+                        y[i0] += sa * sw[i0] * (p[i0] + pr[i0] * (1.0f/254.0f));
+                    }
+                }
+            }
+        }
+    };
+
+    zdnn_tensor_desc view_pre_tfm_desc[6];
+    zdnn_tensor_desc view_tfm_desc[6];
+    std::future<void> sum_task;
+
+    for (int64_t c = 0, g0 = 0; g0 < n_groups; c++, g0 += n_chunk) {
+        const int64_t n = std::min(n_chunk, n_groups - g0);
+
+        const zdnn_ztensor inputs_view     = ggml_zdnn_stack_view(&view_pre_tfm_desc[0], &view_tfm_desc[0], inputs_ztensor,         2 * n_groups, g0,            n);
+        const zdnn_ztensor inputs_res_view = ggml_zdnn_stack_view(&view_pre_tfm_desc[1], &view_tfm_desc[1], inputs_ztensor,         2 * n_groups, n_groups + g0, n);
+        const zdnn_ztensor weights_view    = ggml_zdnn_stack_view(&view_pre_tfm_desc[2], &view_tfm_desc[2], weights_extra->ztensor, n_groups,     g0,            n);
+        const zdnn_ztensor bias_view       = ggml_zdnn_stack_view(&view_pre_tfm_desc[3], &view_tfm_desc[3], bias_ztensor,           n_chunk,      0,             n);
+
+        // the partial buffer of chunk c - 2 is summed by now, as sum_task is waited on before starting the next sum
+        zdnn_ztensor & chunk = partial_ztensor[c % 2];
+        chunk.pre_transformed_desc = &partial_pre_tfm_desc;
+        chunk.transformed_desc     = &partial_tfm_desc;
+        chunk.is_transformed       = false;
+
+        zdnn_ztensor chunk_view     = ggml_zdnn_stack_view(&view_pre_tfm_desc[4], &view_tfm_desc[4], chunk, 2 * n_chunk, 0,       n);
+        zdnn_ztensor chunk_res_view = ggml_zdnn_stack_view(&view_pre_tfm_desc[5], &view_tfm_desc[5], chunk, 2 * n_chunk, n_chunk, n);
+        ZDNN_CHECK(zdnn_quantized_matmul_op(&inputs_view, &weights_view, &bias_view,
+                                            MATMUL_OP_ADDITION, -127, 127, true, false, true,
+                                            nullptr, &chunk_view));
+        ZDNN_CHECK(zdnn_quantized_matmul_op(&inputs_res_view, &weights_view, &bias_view,
+                                            MATMUL_OP_ADDITION, -127, 127, true, false, true,
+                                            nullptr, &chunk_res_view));
+
+        if (sum_task.valid()) {
+            sum_task.get();
+        }
+
+        if (n_threads > 1) {
+            sum_task = std::async(std::launch::async, sum_chunk, chunk, g0, n);
+        } else {
+            sum_chunk(chunk, g0, n);
+        }
+    }
+
+    if (sum_task.valid()) {
+        sum_task.get();
+    }
+
+}
+
+void ggml_zdnn_mul_mat_q(
+          ggml_backend_zdnn_context * ctx,
+    const               ggml_tensor * src0,
+    const               ggml_tensor * src1,
+                        ggml_tensor * dst) {
+    GGML_TENSOR_BINARY_OP_LOCALS;
+
+    const enum ggml_type type = src0->type;
+
+    GGML_ASSERT(ne0 == ne01);
+    GGML_ASSERT(ne1 == ne11);
+    GGML_ASSERT(ne2 == ne12);
+    GGML_ASSERT(ne3 == ne13);
+
+    // we don't support permuted src0 or src1
+    GGML_ASSERT(nb00 == ggml_type_size(type));
+    GGML_ASSERT(nb10 == ggml_type_size(src1->type));
+
+    // dst cannot be transposed or permuted
+    GGML_ASSERT(nb0 == sizeof(float));
+    GGML_ASSERT(nb0 <= nb1);
+    GGML_ASSERT(nb1 <= nb2);
+    GGML_ASSERT(nb2 <= nb3);
+
+    const ggml_tensor * weights = src0;
+    const ggml_tensor * inputs  = src1;
+          ggml_tensor * output  = dst;
+
+    GGML_ASSERT(inputs->type == GGML_TYPE_F32);
+
+    ggml_backend_zdnn_buffer * weights_extra = (ggml_backend_zdnn_buffer *)weights->extra;
+    ggml_backend_zdnn_buffer * output_extra  = (ggml_backend_zdnn_buffer *)output->extra;
+
+    // starting threads costs more than it saves for small inputs, e.g. during token generation
+    const int32_t n_threads = ne10*ne11 >= 65536 ? ctx->n_threads : 1;
+
+    if (weights_extra->pre_tfm_desc.layout == ZDNN_3DS) {
+        GGML_ASSERT(weights_extra->pre_tfm_desc.dim1 == weights->ne[1] && "weights_extra->pre_tfm_desc.dim1 must match weights->ne[1]");
+
+        ggml_zdnn_mul_mat_q_grouped(ctx, n_threads, inputs, weights_extra, output);
+
+        // the result is written to output->data directly, so the output ztensor is stale
+        output_extra->ztensor.is_transformed = false;
+        return;
+    }
+
+    ggml_backend_zdnn_buffer * bias_extra = (ggml_backend_zdnn_buffer *)output_extra->extra;
+
+    GGML_ASSERT(weights_extra->pre_tfm_desc.dim2 == weights->ne[0] && "weights_extra->pre_tfm_desc.dim2 must match weights->ne[0]");
+    GGML_ASSERT(weights_extra->pre_tfm_desc.dim1 == weights->ne[1] && "weights_extra->pre_tfm_desc.dim1 must match weights->ne[1]");
+
+    // zDNN takes one scale per zTensor, so quantize each row of the inputs with its own scale
+    // the rounding errors of the inputs are stacked after the inputs, see ggml_zdnn_quantize_inputs
+    int8_t * inputs_q = ggml_zdnn_resize_scratch_max(ctx->scratch_q, 2 * ne10 * ne11);
+    std::vector<float> inputs_scales(ne11);
+    ggml_zdnn_quantize_inputs(n_threads, inputs, ne10, weights_extra->channel_scales.data(),
+                            inputs_q, inputs_q + ne10 * ne11, inputs_scales.data());
+
+    zdnn_tensor_desc inputs_pre_tfm_desc;
+    zdnn_tensor_desc inputs_tfm_desc;
+    zdnn_ztensor     inputs_ztensor;
+
+    zdnn_init_pre_transformed_desc(ZDNN_3DS, INT8, &inputs_pre_tfm_desc, 2, ne11, ne10);
+    ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&inputs_pre_tfm_desc, QUANTIZED_INT8, &inputs_tfm_desc));
+    ggml_zdnn_init_scratch_ztensor(ctx, &inputs_pre_tfm_desc, &inputs_tfm_desc, &inputs_ztensor, 0);
+    ZDNN_CHECK(zdnn_transform_quantized_ztensor(&inputs_ztensor, false, -127, 127, inputs_q));
+
+    zdnn_tensor_desc partial_pre_tfm_desc;
+    zdnn_tensor_desc partial_tfm_desc;
+    zdnn_ztensor     partial_ztensor;
+
+    zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &partial_pre_tfm_desc, 2, ne11, ne0);
+    ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&partial_pre_tfm_desc, QUANTIZED_DLFLOAT16, &partial_tfm_desc));
+    ggml_zdnn_init_scratch_ztensor(ctx, &partial_pre_tfm_desc, &partial_tfm_desc, &partial_ztensor, 1);
+
+    // both stacks are multiplied by the same unstacked weights in one broadcast call
+    // the bias is all zeros, so it is already the pre-computed bias
+    ZDNN_CHECK(zdnn_quantized_matmul_op(&inputs_ztensor, &weights_extra->ztensor, &bias_extra->ztensor,
+                                        MATMUL_OP_ADDITION, -127, 127, true, false, true,
+                                        nullptr, &partial_ztensor));
+
+    // each stack of the partial results is read back as its own 2D ztensor
+    zdnn_tensor_desc stack_pre_tfm_desc;
+    zdnn_tensor_desc stack_tfm_desc;
+
+    zdnn_init_pre_transformed_desc(ZDNN_2D, FP32, &stack_pre_tfm_desc, ne11, ne0);
+    ZDNN_CHECK(zdnn_generate_transformed_desc(&stack_pre_tfm_desc, &stack_tfm_desc));
+
+    const uint64_t stack_bytes = zdnn_getsize_ztensor(&stack_tfm_desc);
+    GGML_ASSERT(partial_ztensor.buffer_size == 2*stack_bytes);
+
+    float * partial[2] = { (float *)output->data, ggml_zdnn_resize_scratch_max(ctx->scratch_f32, ne11 * ne0) };
+
+    ggml_zdnn_parallel_for(n_threads, 2, [&](int64_t s_start, int64_t s_end) {
+        zdnn_ztensor stack_ztensor;
+        zdnn_init_ztensor(&stack_pre_tfm_desc, &stack_tfm_desc, &stack_ztensor);
+
+        for (int64_t s = s_start; s < s_end; s++) {
+            stack_ztensor.buffer         = (char *)partial_ztensor.buffer + s*stack_bytes;
+            stack_ztensor.buffer_size    = stack_bytes;
+            stack_ztensor.is_transformed = true;
+            ZDNN_CHECK(zdnn_transform_origtensor(&stack_ztensor, partial[s]));
+        }
+    });
+
+    // the result is written to output->data directly, so the output ztensor is stale
+    output_extra->ztensor.is_transformed = false;
+
+    const float * weights_scales = weights_extra->scales.data();
+    ggml_zdnn_parallel_for(n_threads, ne1, [&](int64_t i1_start, int64_t i1_end) {
+        for (int64_t i1 = i1_start; i1 < i1_end; i1++) {
+            const float   sa = inputs_scales[i1];
+            const float * pr = partial[1] + i1*ne0;
+            float       * y  = partial[0] + i1*ne0;
+            for (int64_t i0 = 0; i0 < ne0; i0++) {
+                y[i0] = sa * weights_scales[i0] * (y[i0] + pr[i0] * (1.0f/254.0f));
+            }
+        }
+    });
+}

ggml/src/ggml-zdnn/mmq.hpp

@@ -0,0 +1,12 @@
+#ifndef GGML_ZDNN_MMQ_HPP
+#define GGML_ZDNN_MMQ_HPP
+
+#include "common.hpp"
+
+void ggml_zdnn_mul_mat_q(
+          ggml_backend_zdnn_context * ctx,
+    const               ggml_tensor * src0,
+    const               ggml_tensor * src1,
+                        ggml_tensor * dst);
+
+#endif  // GGML_ZDNN_MMQ_HPP

ggml/src/ggml-zdnn/quantize.cpp

@@ -0,0 +1,58 @@
+#include "ggml.h"
+#include "quantize.hpp"
+
+#include <algorithm>
+#include <cmath>
+
+// quantize each group of group_size values per row to int8, plus the rounding error into residuals for ~15 bits
+void ggml_zdnn_quantize_inputs(int32_t       n_threads,
+                         const ggml_tensor * inputs,
+                               int64_t       group_size,
+                         const float       * channel_scales,
+                               int8_t      * quants,
+                               int8_t      * residuals,
+                               float       * scales) {
+
+    const int64_t ne0      = inputs->ne[0];
+    const int64_t ne1      = inputs->ne[1];
+    const int64_t n_groups = (ne0 + group_size - 1) / group_size;
+
+    ggml_zdnn_parallel_for(n_threads, ne1, [&](int64_t i1_start, int64_t i1_end) {
+        for (int64_t i1 = i1_start; i1 < i1_end; i1++) {
+            const float * xi = (const float *)((const char *)inputs->data + i1*inputs->nb[1]);
+
+            for (int64_t g = 0; g < n_groups; g++) {
+                const float * xg = xi + g*group_size;
+                const float * cg = channel_scales + g*group_size;
+                const int64_t n  = std::min(group_size, ne0 - g*group_size);
+
+                float amax = 0.0f;
+                for (int64_t i0 = 0; i0 < n; i0++) {
+                    amax = std::max(amax, fabsf(xg[i0] * cg[i0]));
+                }
+
+                // the rounding error is at most d/2
+                const float d      = amax / 127.0f;
+                const float id     = d ? 1.0f/d : 0.0f;
+                const float id_res = 254.0f * id;
+                scales[g*ne1 + i1] = d;
+
+                int8_t * qg = quants    + (g*ne1 + i1)*group_size;
+                int8_t * rg = residuals + (g*ne1 + i1)*group_size;
+
+                for (int64_t i0 = 0; i0 < n; i0++) {
+                    const float xc = xg[i0] * cg[i0];
+                    const float v  = roundf(xc * id);
+
+                    qg[i0] = (int8_t)v;
+                    rg[i0] = (int8_t)std::clamp(roundf((xc - v * d) * id_res), -127.0f, 127.0f);
+                }
+
+                for (int64_t i0 = n; i0 < group_size; i0++) {
+                    qg[i0] = 0;
+                    rg[i0] = 0;
+                }
+            }
+        }
+    });
+}

ggml/src/ggml-zdnn/quantize.hpp

@@ -0,0 +1,14 @@
+#ifndef GGML_ZDNN_QUANTIZE_HPP
+#define GGML_ZDNN_QUANTIZE_HPP
+
+#include "common.hpp"
+
+void ggml_zdnn_quantize_inputs(int32_t       n_threads,
+                         const ggml_tensor * inputs,
+                               int64_t       group_size,
+                         const float       * channel_scales,
+                               int8_t      * quants,
+                               int8_t      * residuals,
+                               float       * scales);
+
+#endif  // GGML_ZDNN_QUANTIZE_HPP

ggml/src/ggml-zdnn/utils.cpp

@@ -1,6 +1,10 @@
 #include "ggml.h"
 #include "utils.hpp"
 
+#include <algorithm>
+#include <cmath>
+#include <cstdlib>
+
 zdnn_data_types ggml_zdnn_type_mapping(ggml_type type) {
     switch (type) {
         case GGML_TYPE_F32:
@@ -9,8 +13,6 @@ zdnn_data_types ggml_zdnn_type_mapping(ggml_type type) {
             return FP16;
         case GGML_TYPE_BF16:
             return BFLOAT;
-        case GGML_TYPE_Q8_0:
-            return INT8;
         case GGML_TYPE_I8:
             return INT8;
         case GGML_TYPE_I32:
@@ -43,7 +45,141 @@ void ggml_zdnn_load_tensor(zdnn_ztensor & ztensor, void * buffer) {
     ZDNN_CHECK(zdnn_transform_ztensor(&ztensor, buffer));
 }
 
+// size of the groups along ne0 that get their own scale, or 0 for one scale per row
+// the groups are spread evenly over ne0 and rounded up to 128 values, the last group is zero padded
+int64_t ggml_zdnn_q_group_size(const ggml_tensor * tensor) {
+    // benchmarked on Llama-3.2-1B and Qwen2.5-1.5B Q8_0 with the channel scales:
+    // 2048 is 15-18% faster than 1024, and 1024 lowers the KL divergence to the CPU by 13-19%
+    // GGML_ZDNN_HIGH_PREC=1 selects 1024 for users that want the lower loss
+    // always re-measure with different models to see if different sizes need a dispatching logic
+    static const bool high_prec = std::getenv("GGML_ZDNN_HIGH_PREC") != nullptr &&
+                                  std::atoi(std::getenv("GGML_ZDNN_HIGH_PREC")) != 0;
+    const int64_t group_size_max = high_prec ? 1024 : 2048;
+
+    const int64_t ne0 = tensor->ne[0];
+    if (ne0 <= group_size_max) {
+        return 0;
+    }
+
+    const int64_t n_groups = (ne0 + group_size_max - 1) / group_size_max;
+    return GGML_PAD((ne0 + n_groups - 1) / n_groups, 128);
+}
+
+void ggml_zdnn_load_tensor_q(ggml_backend_zdnn_buffer * buffer, const ggml_tensor * tensor) {
+    const int64_t ne0 = tensor->ne[0];
+    const int64_t ne1 = tensor->ne[1];
+
+    const ggml_to_float_t to_float = ggml_get_type_traits(tensor->type)->to_float;
+    GGML_ASSERT(to_float != nullptr);
+
+    // zDNN takes one scale per zTensor, so requantize each group of each row to int8 with its own scale
+    // a group is the whole row, unless the weights are stacked into groups along ne0
+    // rows are written transposed as zDNN computes inputs (ne11, ne0) x weights (ne0, ne1)
+    const bool    stacked  = buffer->pre_tfm_desc.layout == ZDNN_3DS;
+    const int64_t n_groups = stacked ? buffer->pre_tfm_desc.dim3 : 1;
+    const int64_t gs       = stacked ? buffer->pre_tfm_desc.dim2 : ne0;
+
+    // benchmarked: 256 rows fill a 256-byte cache line per write, padding each row by a line avoids cache set conflicts
+    constexpr int64_t n_tile   = 256;
+    const     int64_t x_stride = ne0 + 64;
+
+    // the rows past ne0 of the last group stay zero as padding
+    std::vector<int8_t> q(n_groups * gs * ne1);
+    std::vector<float>  x(n_tile * x_stride);
+    std::vector<float>  id(n_tile * n_groups);
+
+    // a few input channels with larger weights than the rest would set the scale of every group,
+    // so divide each input channel by its RMS over the rows and multiply the inputs by it instead
+    // with groups of 1024, this lowered the KL divergence to the CPU by 10-18% on Llama-3.2-1B and Qwen2.5-1.5B Q8_0 at the same speed
+    std::vector<double> sum_sq(ne0, 0.0);
+    for (int64_t i1 = 0; i1 < ne1; i1++) {
+        to_float((const char *)tensor->data + i1*tensor->nb[1], x.data(), ne0);
+        for (int64_t i0 = 0; i0 < ne0; i0++) {
+            sum_sq[i0] += (double)x[i0] * x[i0];
+        }
+    }
+
+    std::vector<float> channel_id(ne0);
+    for (int64_t i0 = 0; i0 < ne0; i0++) {
+        const float c = (float)std::sqrt(sum_sq[i0] / ne1);
+        buffer->channel_scales[i0] = c ? c : 1.0f;
+        channel_id[i0] = 1.0f / buffer->channel_scales[i0];
+    }
+
+    for (int64_t i1 = 0; i1 < ne1; i1 += n_tile) {
+        const int64_t nt = std::min(n_tile, ne1 - i1);
+
+        for (int64_t t = 0; t < nt; t++) {
+            float * xt = x.data() + t*x_stride;
+            to_float((const char *)tensor->data + (i1 + t)*tensor->nb[1], xt, ne0);
+
+            for (int64_t i0 = 0; i0 < ne0; i0++) {
+                xt[i0] *= channel_id[i0];
+            }
+
+            for (int64_t g = 0; g < n_groups; g++) {
+                float amax = 0.0f;
+                for (int64_t i0 = g*gs; i0 < std::min((g + 1)*gs, ne0); i0++) {
+                    amax = std::max(amax, fabsf(xt[i0]));
+                }
+
+                const float d = amax / 127.0f;
+                buffer->scales[g*ne1 + i1 + t] = d;
+                id[t*n_groups + g] = d ? 1.0f/d : 0.0f;
+            }
+        }
+
+        for (int64_t i0 = 0; i0 < ne0; i0++) {
+            const int64_t g = i0 / gs;
+            int8_t * qi = q.data() + i0*ne1 + i1;
+            for (int64_t t = 0; t < nt; t++) {
+                qi[t] = (int8_t)roundf(x[t*x_stride + i0] * id[t*n_groups + g]);
+            }
+        }
+    }
+
+    if (buffer->ztensor.is_transformed) {
+        zdnn_reset_ztensor(&buffer->ztensor);
+    }
+    ZDNN_CHECK(zdnn_transform_quantized_ztensor(&buffer->ztensor, false, -127, 127, q.data()));
+}
+
 void ggml_zdnn_init_tensor(ggml_backend_zdnn_buffer * buffer, const ggml_tensor * tensor) {
+    if (ggml_is_quantized(tensor->type)) {
+        GGML_ASSERT(ggml_is_matrix(tensor));
+
+        const int64_t gs = ggml_zdnn_q_group_size(tensor);
+
+        if (gs > 0) {
+            // stack the weights into groups along ne0, each group of each row gets its own scale
+            const int64_t n_groups = (tensor->ne[0] + gs - 1) / gs;
+
+            zdnn_init_pre_transformed_desc(
+                ZDNN_3DS,
+                INT8,
+                &buffer->pre_tfm_desc,
+                n_groups, gs, tensor->ne[1]
+            );
+
+            buffer->scales.resize(n_groups * tensor->ne[1]);
+        } else {
+            zdnn_init_pre_transformed_desc(
+                ZDNN_2D,
+                INT8,
+                &buffer->pre_tfm_desc,
+                tensor->ne[0], tensor->ne[1]
+            );
+
+            buffer->scales.resize(tensor->ne[1]);
+        }
+
+        buffer->channel_scales.resize(tensor->ne[0]);
+
+        ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&buffer->pre_tfm_desc, QUANTIZED_WEIGHTS_INT8, &buffer->tfm_desc));
+        ZDNN_CHECK(zdnn_init_quantized_ztensor_with_malloc(&buffer->pre_tfm_desc, &buffer->tfm_desc, 1.0f, 0.0f, &buffer->ztensor));
+        return;
+    }
+
     switch (tensor->op) {
         case GGML_OP_MUL_MAT:
             {

ggml/src/ggml-zdnn/utils.hpp

@@ -14,6 +14,10 @@ void ggml_zdnn_create_tensor(zdnn_tensor_desc & pre_tfm_desc,
 
 void ggml_zdnn_load_tensor(zdnn_ztensor & ztensor, void * buffer);
 
+int64_t ggml_zdnn_q_group_size(const ggml_tensor * tensor);
+
+void ggml_zdnn_load_tensor_q(ggml_backend_zdnn_buffer * buffer, const ggml_tensor * tensor);
+
 void ggml_zdnn_init_tensor(ggml_backend_zdnn_buffer * buffer, const ggml_tensor * tensor);
 
 #endif  // GGML_ZDNN_UTILITIES_HPP

ggml/src/ggml.c

@@ -329,6 +329,17 @@ void ggml_log_callback_default(enum ggml_log_level level, const char * text, voi
 #endif
 
 
+// the thread sanitizer does not intercept vm_allocate and vm_deallocate, so it produces false positives
+#if defined(__has_feature)
+#if __has_feature(thread_sanitizer)
+#define GGML_TSAN_ENABLED 1
+#endif
+#else
+#if defined(__SANITIZE_THREAD__)
+#define GGML_TSAN_ENABLED 1
+#endif
+#endif
+
 void * ggml_aligned_malloc(size_t size) {
 #if defined(__s390x__)
     const int alignment = 256;
@@ -346,7 +357,7 @@ void * ggml_aligned_malloc(size_t size) {
     void * aligned_memory = NULL;
   #ifdef GGML_USE_CPU_HBM
     int result = hbw_posix_memalign(&aligned_memory, alignment, size);
-  #elif TARGET_OS_OSX
+  #elif TARGET_OS_OSX && !defined(GGML_TSAN_ENABLED)
     GGML_UNUSED(alignment);
     kern_return_t alloc_status = vm_allocate((vm_map_t) mach_task_self(), (vm_address_t *) &aligned_memory, size, VM_FLAGS_ANYWHERE);
     int result = EFAULT;
@@ -393,7 +404,7 @@ void ggml_aligned_free(void * ptr, size_t size) {
     if (ptr != NULL) {
         hbw_free(ptr);
     }
-#elif TARGET_OS_OSX
+#elif TARGET_OS_OSX && !defined(GGML_TSAN_ENABLED)
     if (ptr != NULL) {
         vm_deallocate((vm_map_t)mach_task_self(), (vm_address_t)ptr, size);
     }

include/llama.h

@@ -405,7 +405,7 @@ extern "C" {
         void *              abort_callback_data;
 
         // Keep the booleans together and at the end of the struct to avoid misalignment during copy-by-value.
-        bool embeddings;  // if true, extract embeddings (together with logits)
+        bool embeddings;  // if true, extract embeddings (together with logits, except for a pooling type other than NONE)
         bool offload_kqv; // offload the KQV ops (including the KV cache) to GPU
         bool no_perf;     // measure performance timings
         bool op_offload;  // offload host tensor operations to device

src/llama-context.cpp

@@ -789,6 +789,8 @@ void llama_context::sched_reserve() {
         n_nodes_pp         = ggml_graph_n_nodes(gf);
         n_inputs_pp        = get_gf_res_reserve()->inputs.size();
         n_input_tensors_pp = this->n_input_tensors;
+
+        graph_has_logits = get_gf_res_reserve()->get_logits() != nullptr;
     }
 
     // reserve with tg (token generation) graph to get the number of splits and nodes
@@ -1321,10 +1323,12 @@ void llama_context::set_abort_callback(bool (*abort_callback)(void * data), void
 void llama_context::set_embeddings(bool value) {
     LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value);
 
-    cparams.embeddings = value;
+    // a pooled context computes the logits only without embeddings, so the toggle changes the graph shape
+    if (cparams.embeddings != value && cparams.pooling_type != LLAMA_POOLING_TYPE_NONE) {
+        sched_need_reserve = true;
+    }
 
-    // TODO: not sure yet if we want to reserve here
-    //sched_need_reserve = true;
+    cparams.embeddings = value;
 }
 
 void llama_context::set_embeddings_nextn(bool value, bool masked) {
@@ -2231,7 +2235,7 @@ uint32_t llama_context::output_reserve(int32_t n_outputs) {
     const auto n_embd     = hparams.n_embd;
     const auto n_embd_out = hparams.n_embd_out();
 
-    bool has_logits     = true;
+    bool has_logits     = !cparams.embeddings || graph_has_logits;
     bool has_embd       = cparams.embeddings;
     bool has_embd_nextn = cparams.embeddings_nextn;
 
@@ -2719,8 +2723,8 @@ void llama_context::sched_graph_realloc(ggml_backend_sched_t sched, ggml_cgraph
         lctx->n_reallocs_same++;
     }
 
-    LLAMA_LOG_INFO("%s: graph reallocation: nodes = %d, splits = %d, n_same = %d, n_diff = %d\n",
-            __func__, ggml_graph_n_nodes(gf), ggml_backend_sched_get_n_splits(sched),
+    LLAMA_LOG_INFO("%s: graph reallocation: nodes = %d (changed = %d), splits = %d, n_same = %d, n_diff = %d\n",
+            __func__, ggml_graph_n_nodes(gf), graph_size_changed, ggml_backend_sched_get_n_splits(sched),
             lctx->n_reallocs_same, lctx->n_reallocs_diff);
 }
 

src/llama-context.h

@@ -280,7 +280,7 @@ private:
     static bool sched_copy_experts(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph, void * user_data);
 
     // ggml_backend_sched realloc callback, counts and traces the graph reallocations of this context
-    static void sched_graph_realloc(ggml_backend_sched_t sched, ggml_cgraph * graph, bool backend_ids_changed, bool graph_size_changed, void * user_data);
+    static void sched_graph_realloc(ggml_backend_sched_t sched, ggml_cgraph * gf, bool backend_ids_changed, bool graph_size_changed, void * user_data);
 
     // disable auto fused ops (Flash Attention, Gated Delta Net) whose op lands on a device
     // that differs from the layer it belongs to (usually due to missing backend support)
@@ -368,6 +368,10 @@ private:
 
     bool sched_need_reserve = true;
 
+    // whether the reserve graph has a logits tensor, set in sched_reserve()
+    // true until then, because the first output_reserve() runs before any graph is reserved
+    bool graph_has_logits = true;
+
     // state of sched_copy_experts, reset before each graph compute
     struct copy_experts_info {
         const ggml_tensor *  ids = nullptr;

src/llama-ext.h

@@ -8,6 +8,7 @@
 
 #include <cstdint>
 #include <map>
+#include <vector>
 
 // Reserve a new compute graph. It is valid until the next call to llama_graph_reserve.
 LLAMA_API struct ggml_cgraph * llama_graph_reserve(
@@ -90,6 +91,13 @@ LLAMA_API ggml_backend_dev_t llama_model_get_device(const struct llama_model * m
 
 LLAMA_API llama_memory_breakdown llama_get_memory_breakdown(const struct llama_context * ctx);
 
+// the backend buffers that a state write of this memory would touch - empty caches report none
+LLAMA_API std::vector<ggml_backend_buffer_t> llama_memory_get_buffers(const llama_memory_t mem);
+
+// Turn the KV cache attention rotation on or off for the contexts created from this model.
+// Useful for testing purposes.
+LLAMA_API void llama_model_set_attn_rot(struct llama_model * model, bool value);
+
 // Set whether the context outputs nextn embeddings or not
 // If masked == true,  output the embeddings only for the tokens with batch.logits != 0
 // If masked == false, output the embeddings for all tokens in the batch regardless of batch.logits

src/llama-graph.cpp

@@ -3058,10 +3058,9 @@ ggml_tensor * llm_graph_context::build_attn(
 
     // the KV values are padded with zeros up to the KV dimension of the mask
     const int64_t n_kv_pad = kq_mask->ne[0] - k->ne[2];
-    if (n_kv_pad > 0) {
-        k = ggml_pad(ctx0, k, 0, 0, n_kv_pad, 0);
-        v = ggml_pad(ctx0, v, 0, 0, n_kv_pad, 0);
-    }
+
+    k = ggml_pad(ctx0, k, 0, 0, n_kv_pad, 0);
+    v = ggml_pad(ctx0, v, 0, 0, n_kv_pad, 0);
 
     ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, 0, kq_scale, il);
     cb(cur, "kqv_out", il);
@@ -4064,6 +4063,18 @@ void llm_graph_context::build_pooling(
     ggml_build_forward_expand(gf, cur);
 }
 
+void llm_graph_context::build_logits() const {
+    // a pooled embeddings context outputs one vector per sequence, so its graph skips the LM head
+    if (cparams.embeddings && pooling_type != LLAMA_POOLING_TYPE_NONE) {
+        res->t_logits = nullptr;
+    }
+
+    // the models declare t_logits without expanding it, so the LM head is computed only when it is expanded here
+    if (res->t_logits != nullptr) {
+        ggml_build_forward_expand(gf, res->t_logits);
+    }
+}
+
 void llm_graph_context::build_sampling() const {
     if (samplers.empty() || !res->t_logits) {
         return;

src/llama-graph.h

@@ -1429,6 +1429,12 @@ struct llm_graph_context {
             ggml_tensor * cls_out_b,
             ggml_tensor * cls_norm) const;
 
+    //
+    // logits
+    //
+
+    void build_logits() const;
+
     //
     // sampling (backend sampling)
     //

src/llama-kv-cache.cpp

@@ -312,10 +312,14 @@ llama_kv_cache::llama_kv_cache(
         n_rot_k = other->n_rot_k;
         n_rot_v = other->n_rot_v;
     } else {
-        const char * LLAMA_ATTN_ROT_DISABLE = getenv("LLAMA_ATTN_ROT_DISABLE");
-        const bool attn_rot_disable = LLAMA_ATTN_ROT_DISABLE ? atoi(LLAMA_ATTN_ROT_DISABLE) : false;
+        // two knobs turn the same thing off: the environment variable, for the whole process, and
+        //   the model setting, for the contexts of one model - see llama_model_set_attn_rot
+        const char * env_rot_disable = getenv("LLAMA_ATTN_ROT_DISABLE");
+
+        const bool attn_rot_disable = !model.attn_rot || (env_rot_disable != nullptr && atoi(env_rot_disable) != 0);
         if (attn_rot_disable) {
-            LLAMA_LOG_WARN("%s: attention rotation force disabled (LLAMA_ATTN_ROT_DISABLE)\n", __func__);
+            LLAMA_LOG_WARN("%s: attention rotation force disabled (%s)\n", __func__,
+                    model.attn_rot ? "LLAMA_ATTN_ROT_DISABLE" : "llama_model_set_attn_rot");
         }
 
         bool attn_rot_k =

src/llama-memory.cpp

@@ -1,5 +1,11 @@
 #include "llama-memory.h"
 
+#include "llama-ext.h"
+#include "llama-io.h"
+
+#include <set>
+#include <vector>
+
 llama_memory_status llama_memory_status_combine(llama_memory_status s0, llama_memory_status s1) {
     bool has_update = false;
 
@@ -57,3 +63,30 @@ bool llama_memory_status_is_fail(llama_memory_status status) {
 
     return false;
 }
+
+// collects the buffers of the tensors that a state write visits, without keeping the data
+struct llama_memory_buffers_writer : public llama_io_write_i {
+    std::set<ggml_backend_buffer_t> buffers;
+    size_t size = 0;
+
+    void write(const void *, size_t n) override {
+        size += n;
+    }
+
+    void write_tensor(ggml_tensor * tensor, size_t, size_t n) override {
+        buffers.insert(tensor->buffer);
+        size += n;
+    }
+
+    size_t n_bytes() override {
+        return size;
+    }
+};
+
+std::vector<ggml_backend_buffer_t> llama_memory_get_buffers(const llama_memory_t mem) {
+    llama_memory_buffers_writer writer;
+
+    mem->state_write(writer);
+
+    return std::vector<ggml_backend_buffer_t>(writer.buffers.begin(), writer.buffers.end());
+}

src/llama-mmap.cpp

@@ -835,7 +835,7 @@ const bool llama_mlock::SUPPORTED = false;
 #endif
 
 void llama_prefetch(llama_memory_ranges mr) {
-#if defined(__linux__) || (defined(_WIN32) && _WIN32_WINNT >= 0x602)
+#if defined(__linux__) || defined(__APPLE__) || (defined(_WIN32) && _WIN32_WINNT >= 0x602)
     if (mr.empty()) {
         return;
     }
@@ -874,6 +874,7 @@ void llama_prefetch(llama_memory_ranges mr) {
     };
 #else
     auto prefetch = [&]() {
+        // TODO: in a follow-up, use posix_madvise(.., POSIX_MADV_WILLNEED) to match the rest of this file
         if (madvise((void *) begin, end - begin, MADV_WILLNEED) != 0) {
             LLAMA_LOG_WARN("llama_prefetch: madvise(MADV_WILLNEED) failed: %s\n", strerror(errno));
             return false;

src/llama-model-loader.cpp

@@ -13,6 +13,7 @@
 #include <cstring>
 #include <future>
 #include <regex>
+#include <unordered_set>
 
 static const size_t kiB = 1024;
 static const size_t MiB = 1024*kiB;
@@ -1446,21 +1447,80 @@ void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps
     }
 }
 
-void llama_model_loader::get_mapping_range(size_t * first, size_t * last, void ** addr, int idx, ggml_context * ctx) const {
-    GGML_ASSERT(!mappings.empty());
-    const auto & mapping = mappings.at(idx);
+using mapping_ranges = std::vector<llama_model_loader::mapping_range>;
+
+// merge a list of [first, last) ranges into a sorted, non-overlapping list
+static mapping_ranges merge_ranges(mapping_ranges ranges) {
+    std::sort(ranges.begin(), ranges.end(), [](const auto & a, const auto & b) { return a.first < b.first; });
+
+    mapping_ranges res;
+    for (const auto & range : ranges) {
+        if (!res.empty() && range.first <= res.back().last) {
+            res.back().last = std::max(res.back().last, range.last);
+        } else {
+            res.push_back(range);
+        }
+    }
+
+    return res;
+}
+
+// subtract the sorted, non-overlapping ranges 'sub' from 'a'
+static mapping_ranges sub_ranges(const mapping_ranges & a, const mapping_ranges & sub) {
+    mapping_ranges res;
+
+    for (const auto & x : a) {
+        size_t cur = x.first;
+        for (const auto & y : sub) {
+            if (y.last <= cur || y.first >= x.last) {
+                continue;
+            }
+            if (y.first > cur) {
+                res.push_back({ cur, y.first });
+            }
+            cur = std::max(cur, y.last);
+        }
+        if (cur < x.last) {
+            res.push_back({ cur, x.last });
+        }
+    }
+
+    return res;
+}
+
+std::vector<llama_model_loader::mapping_range> llama_model_loader::get_mapping_ranges(uint32_t idx, ggml_context * ctx) const {
+    GGML_ASSERT(idx < mappings.size());
+
+    // data of this context, and data of the other contexts that share the same file
+    // a weight that is mapped in several contexts (e.g. the token embedding) is not foreign data
+    std::unordered_set<const llama_tensor_weight *> own;
+    mapping_ranges own_ranges;
+    mapping_ranges other;
 
-    *first = mapping->size();
-    *last  = 0;
-    *addr = mapping->addr();
     for (ggml_tensor * tensor = ggml_get_first_tensor(ctx); tensor; tensor = ggml_get_next_tensor(ctx, tensor)) {
         const auto * weight = get_weight(ggml_get_name(tensor));
         if (!weight || weight->idx != idx) {
             continue;
         }
-        *first = std::min(*first, weight->offs);
-        *last  = std::max(*last,  weight->offs + ggml_nbytes(tensor));
+        own.insert(weight);
+        own_ranges.push_back({ weight->offs, weight->offs + ggml_nbytes(tensor) });
     }
+
+    for (const auto & it : ctx_map) {
+        if (!it.second || it.second.get() == ctx) {
+            continue;
+        }
+
+        for (ggml_tensor * tensor = ggml_get_first_tensor(it.second.get()); tensor; tensor = ggml_get_next_tensor(it.second.get(), tensor)) {
+            const auto * weight = get_weight(ggml_get_name(tensor));
+            if (!weight || weight->idx != idx || own.count(weight)) {
+                continue;
+            }
+            other.push_back({ weight->offs, weight->offs + ggml_nbytes(tensor) });
+        }
+    }
+
+    return sub_ranges(merge_ranges(std::move(own_ranges)), merge_ranges(std::move(other)));
 }
 
 void llama_model_loader::unmap_weight(const llama_tensor_weight & w) const {
@@ -1492,7 +1552,7 @@ const void * llama_model_loader::load_data_range(const llama_tensor_weight & w,
 
 bool llama_model_loader::load_all_data(
         struct ggml_context * ctx,
-        llama_buf_map & bufs,
+        std::vector<ggml_backend_buffer_t> & bufs,
         llama_mlocks * lmlocks,
         llama_progress_callback progress_callback,
         void * progress_callback_user_data) {
@@ -1528,7 +1588,7 @@ bool llama_model_loader::load_all_data(
         }
         // When not using mmaped io use async uploads from pinned memory to GPU memory.
         // First determine if the backend supports the necessary features for async uploads.
-        auto * buf = bufs.count(0) ? bufs.at(0) : nullptr;
+        auto * buf = bufs.empty() ? nullptr : bufs[0];
         if (!buf) {
             LLAMA_LOG_DEBUG("%s: no buffer found for async uploads\n", func);
             return nullptr;
@@ -1599,7 +1659,7 @@ bool llama_model_loader::load_all_data(
     if (upload_backend) {
         LLAMA_LOG_DEBUG("%s: using async uploads for device %s, buffer type %s, backend %s\n", __func__,
             ggml_backend_dev_name(ggml_backend_get_device(upload_backend)),
-            ggml_backend_buft_name(ggml_backend_buffer_get_type(bufs.at(0))),
+            ggml_backend_buft_name(ggml_backend_buffer_get_type(bufs[0])),
             ggml_backend_name(upload_backend));
     }
 
@@ -1640,12 +1700,18 @@ bool llama_model_loader::load_all_data(
 
         if (from_mapping) {
             const auto & mapping = mappings.at(weight->idx);
-            ggml_backend_buffer_t buf_mmap = nullptr;
-            if (bufs.count(weight->idx)) {
-                buf_mmap = bufs.at(weight->idx);
-            }
             uint8_t * data = (uint8_t *) mapping->addr() + weight->offs;
 
+            // the tensors of a file can be spread across several buffers - find the one containing this tensor's data
+            ggml_backend_buffer_t buf_mmap = nullptr;
+            for (const auto & buf : bufs) {
+                const uint8_t * base = (const uint8_t *) ggml_backend_buffer_get_base(buf);
+                if (base && base <= data && data + n_size <= base + ggml_backend_buffer_get_size(buf)) {
+                    buf_mmap = buf;
+                    break;
+                }
+            }
+
             if (check_tensors) {
                 validation_result.emplace_back(std::async(std::launch::async, [cur, data, n_size] {
                     return std::make_pair(cur, ggml_validate_row_data(cur->type, data, n_size));

src/llama-model-loader.h

@@ -16,8 +16,6 @@
 #include <stdexcept>
 #include <unordered_map>
 
-using llama_buf_map = std::unordered_map<uint32_t, ggml_backend_buffer_t>;
-
 // lists of buffer types used for each layer
 using buft_list_t = std::vector<std::pair<ggml_backend_dev_t, ggml_backend_buffer_type_t>>;
 
@@ -242,7 +240,16 @@ struct llama_model_loader {
 
     void init_mappings(bool prefetch = true, llama_mlocks * mlock_mmaps = nullptr);
 
-    void get_mapping_range(size_t * first, size_t * last, void ** addr, int idx, ggml_context * ctx) const;
+    // byte range within a mapped file
+    struct mapping_range {
+        size_t first; // inclusive
+        size_t last;  // exclusive
+    };
+
+    // sorted, non-overlapping ranges of file 'idx' that back the tensors of 'ctx'
+    // ranges owned by other contexts are excluded, so that a buffer wrapping them won't pin
+    // data that this context never reads (lazy tensors, -cmoe tensors, ...)
+    std::vector<mapping_range> get_mapping_ranges(uint32_t idx, ggml_context * ctx) const;
 
     // release a weight's mmap pages
     void unmap_weight(const llama_tensor_weight & w) const;
@@ -254,7 +261,7 @@ struct llama_model_loader {
     // Returns false if cancelled by progress_callback
     bool load_all_data(
             struct ggml_context * ctx,
-            llama_buf_map & bufs,
+            std::vector<ggml_backend_buffer_t> & bufs,
             llama_mlocks * lmlocks,
             llama_progress_callback progress_callback,
             void * progress_callback_user_data);

src/llama-model.cpp

@@ -1875,12 +1875,13 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
     pimpl->mappings.reserve(ml.mappings.size());
 
     // create the backend buffers
-    std::vector<std::pair<ggml_context *, llama_buf_map>> ctx_buf_maps;
-    ctx_buf_maps.reserve(ml.ctx_map.size());
+    // note: the buffers are owned by pimpl->ctxs_bufs, this is a raw-pointer view used to load the data
+    std::vector<std::pair<ggml_context *, std::vector<ggml_backend_buffer_t>>> ctx_bufs;
+    ctx_bufs.reserve(ml.ctx_map.size());
 
-    // Ensure we have enough capacity for the maximum backend buffer we will potentially create
-    const size_t n_max_backend_buffer = ml.ctx_map.size() * ml.files.size();
-    pimpl->ctxs_bufs.reserve(n_max_backend_buffer);
+    // reserve hint: each context creates one buffer per file, or more when its tensors are spread
+    //     across disjoint ranges of a file (see get_mapping_ranges)
+    pimpl->ctxs_bufs.reserve(ml.ctx_map.size() * ml.files.size());
 
     for (auto & [ctx_key, ctx_ptr] : ml.ctx_map) {
         ggml_backend_buffer_type_t buft = ctx_key.buft;
@@ -1891,9 +1892,6 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
             continue;
         }
 
-        llama_buf_map buf_map;
-        buf_map.reserve(n_max_backend_buffer);
-
         // check if it is possible to use buffer_from_host_ptr with this buffer type
         ggml_backend_dev_t dev = ggml_backend_buft_get_device(buft);
         if (!dev) {
@@ -1915,24 +1913,24 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
 
         if ((ml.use_mmap || is_lazy_mapped) && use_mmap_buffer && buffer_from_host_ptr_supported && is_default_buft) {
             GGML_ASSERT(!ml.no_alloc);
+            const size_t max_size = ggml_get_max_tensor_size(ctx);
+
             for (uint32_t idx = 0; idx < ml.files.size(); idx++) {
-                // only the mmap region containing the tensors in the model is mapped to the backend buffer
-                // this is important for metal with apple silicon: if the entire model could be mapped to a metal buffer,
-                //     then we could just use metal for all layers
-                // this allows using partial offloading when the model size exceeds the metal buffer size, but not the RAM size
-                void * addr = nullptr;
-                size_t first, last; // NOLINT
-                ml.get_mapping_range(&first, &last, &addr, idx, ctx);
-                if (first >= last) {
-                    continue;
+                // map the mmap regions that contain the tensors of this context
+                // this is important for metal with apple silicon: a backend buffer pins everything it
+                //     wraps, so the regions owned by the other contexts (lazy tensors read on demand
+                //     from disk, MoE tensors kept on the CPU, etc) are left out and stay out of RAM
+                // this also allows using partial offloading when the model size exceeds the metal buffer
+                //     size, but not the RAM size
+                void * addr = ml.mappings.at(idx)->addr();
+
+                for (const auto & range : ml.get_mapping_ranges(idx, ctx)) {
+                    ggml_backend_buffer_t buf = ggml_backend_dev_buffer_from_host_ptr(dev, (char *) addr + range.first, range.last - range.first, max_size);
+                    if (buf == nullptr) {
+                        throw std::runtime_error(format("unable to allocate %s buffer", ggml_backend_buft_name(buft)));
+                    }
+                    bufs.emplace_back(buf);
                 }
-                const size_t max_size = ggml_get_max_tensor_size(ctx);
-                ggml_backend_buffer_t buf = ggml_backend_dev_buffer_from_host_ptr(dev, (char *) addr + first, last - first, max_size);
-                if (buf == nullptr) {
-                    throw std::runtime_error(format("unable to allocate %s buffer", ggml_backend_buft_name(buft)));
-                }
-                bufs.emplace_back(buf);
-                buf_map.emplace(idx, buf);
             }
         } else {
             ggml_backend_buffer_t buf;
@@ -1954,20 +1952,22 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
                 mlock_buf->grow_to(ggml_backend_buffer_get_size(buf));
             }
             bufs.emplace_back(buf);
-            for (uint32_t idx = 0; idx < ml.files.size(); idx++) {
-                buf_map.emplace(idx, buf);
-            }
         }
 
+        std::vector<ggml_backend_buffer_t> buf_ptrs;
+        buf_ptrs.reserve(bufs.size());
+
         for (auto & buf : bufs) {
             // indicate that this buffer contains weights
             // this is used by ggml_backend_sched to improve op scheduling: ops that use a weight are preferably scheduled to the backend that contains the weight
             ggml_backend_buffer_set_usage(buf.get(), GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
+
+            buf_ptrs.emplace_back(buf.get());
         }
 
         pimpl->ctxs_bufs.emplace_back(std::move(ctx_ptr), std::move(bufs));
 
-        ctx_buf_maps.emplace_back(ctx, buf_map);
+        ctx_bufs.emplace_back(ctx, std::move(buf_ptrs));
     }
 
     if (llama_supports_gpu_offload()) {
@@ -2000,15 +2000,15 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
 
     // without mmap, load non-host buffers first: their tensors go through a staging buffer, which is cheapest while the fewest weights are resident
     if (!ml.use_mmap) {
-        std::stable_partition(ctx_buf_maps.begin(), ctx_buf_maps.end(), [](const auto & ctx_buf_map) {
-            const auto & buf_map = ctx_buf_map.second;
-            return !buf_map.empty() && !ggml_backend_buffer_is_host(buf_map.begin()->second);
+        std::stable_partition(ctx_bufs.begin(), ctx_bufs.end(), [](const auto & x) {
+            const auto & bufs = x.second;
+            return !bufs.empty() && !ggml_backend_buffer_is_host(bufs[0]);
         });
     }
 
     // load tensor data
-    for (auto & [ctx, buf_map] : ctx_buf_maps) {
-        if (!ml.load_all_data(ctx, buf_map, use_mlock ? &pimpl->mlock_mmaps : NULL, params.progress_callback, params.progress_callback_user_data)) {
+    for (auto & [ctx, bufs] : ctx_bufs) {
+        if (!ml.load_all_data(ctx, bufs, use_mlock ? &pimpl->mlock_mmaps : NULL, params.progress_callback, params.progress_callback_user_data)) {
             return false;
         }
     }
@@ -3305,6 +3305,9 @@ ggml_cgraph * llama_model::build_graph(const llm_graph_params & params) const {
     // add on pooling layer
     llm->build_pooling(cls, cls_b, cls_out, cls_out_b, cls_norm);
 
+    // add the LM head (if the context outputs logits)
+    llm->build_logits();
+
     // add backend sampling layers (if any)
     llm->build_sampling();
 
@@ -3781,6 +3784,10 @@ ggml_backend_dev_t llama_model_get_device(const struct llama_model * model, int
     return model->devices[i].dev;
 }
 
+void llama_model_set_attn_rot(struct llama_model * model, bool value) {
+    model->attn_rot = value;
+}
+
 //
 // llama_model_base
 //

src/llama-model.h

@@ -690,6 +690,9 @@ struct llama_model {
 
     llama_hadamard hdmd;
 
+    // see llama_model_set_attn_rot
+    bool attn_rot = true;
+
     // for classifier models
     std::vector<std::string> classifier_labels;
 

src/models/afmoe.cpp

@@ -278,6 +278,4 @@ llama_model_afmoe::graph::graph(const llama_model & model, const llm_graph_param
     cur = build_lora_mm(model.output, cur, model.output_s);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/apertus.cpp

@@ -165,6 +165,4 @@ llama_model_apertus::graph::graph(const llama_model & model, const llm_graph_par
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/arcee.cpp

@@ -152,6 +152,4 @@ llama_model_arcee::graph::graph(const llama_model & model, const llm_graph_param
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/arctic.cpp

@@ -175,6 +175,4 @@ llama_model_arctic::graph::graph(const llama_model & model, const llm_graph_para
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/arwkv7.cpp

@@ -197,6 +197,4 @@ llama_model_arwkv7::graph::graph(const llama_model & model, const llm_graph_para
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/baichuan.cpp

@@ -150,6 +150,4 @@ llama_model_baichuan::graph::graph(const llama_model & model, const llm_graph_pa
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/bailingmoe.cpp

@@ -175,6 +175,4 @@ llama_model_bailingmoe::graph::graph(const llama_model & model, const llm_graph_
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/bailingmoe2.cpp

@@ -206,6 +206,4 @@ llama_model_bailingmoe2::graph::graph(const llama_model & model, const llm_graph
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/bailingmoe3.cpp

@@ -412,7 +412,6 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph
     cur = ggml_mul_mat(ctx0, model.output, cur);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-    ggml_build_forward_expand(gf, cur);
 }
 
 llama_model_bailingmoe3::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) :
@@ -555,5 +554,4 @@ llama_model_bailingmoe3::graph_mtp::graph_mtp(const llama_model & model, const l
     cur = ggml_mul_mat(ctx0, model.output, cur);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-    ggml_build_forward_expand(gf, cur);
 }

src/models/bert.cpp

@@ -73,11 +73,17 @@ llama_model_bert::graph::graph(const llama_model & model, const llm_graph_params
     ggml_tensor * cur;
     ggml_tensor * inpL;
     ggml_tensor * inp_pos = nullptr;
+    ggml_tensor * embd_pos = nullptr;
 
     if (model.arch != LLM_ARCH_JINA_BERT_V2) {
         inp_pos = build_inp_pos();
     }
 
+    if (model.arch == LLM_ARCH_BERT) {
+        embd_pos = ggml_get_rows(ctx0, model.pos_embd, inp_pos);
+        ggml_build_forward_order(gf, embd_pos);
+    }
+
     // construct input embeddings (token, type, position)
     inpL = build_inp_embd(model.tok_embd);
 
@@ -87,7 +93,7 @@ llama_model_bert::graph::graph(const llama_model & model, const llm_graph_params
         inpL                    = ggml_add(ctx0, inpL, type_row0);
     }
     if (model.arch == LLM_ARCH_BERT) {
-        inpL = ggml_add(ctx0, ggml_get_rows(ctx0, model.pos_embd, inp_pos), inpL);
+        inpL = ggml_add(ctx0, embd_pos, inpL);
     }
     cb(inpL, "inp_embd", -1);
 

src/models/bitnet.cpp

@@ -165,6 +165,4 @@ llama_model_bitnet::graph::graph(const llama_model & model, const llm_graph_para
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/bloom.cpp

@@ -146,6 +146,4 @@ llama_model_bloom::graph::graph(const llama_model & model, const llm_graph_param
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/chameleon.cpp

@@ -199,6 +199,4 @@ llama_model_chameleon::graph::graph(const llama_model & model, const llm_graph_p
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/chatglm.cpp

@@ -156,6 +156,4 @@ llama_model_chatglm::graph::graph(const llama_model & model, const llm_graph_par
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/codeshell.cpp

@@ -148,6 +148,4 @@ llama_model_codeshell::graph::graph(const llama_model & model, const llm_graph_p
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/cogvlm.cpp

@@ -155,5 +155,4 @@ llama_model_cogvlm::graph::graph(const llama_model & model, const llm_graph_para
     cur = build_lora_mm(model.output, cur, model.output_s);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-    ggml_build_forward_expand(gf, cur);
 }

src/models/cohere2.cpp

@@ -154,6 +154,4 @@ llama_model_cohere2::graph::graph(const llama_model & model, const llm_graph_par
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/cohere2moe.cpp

@@ -275,8 +275,6 @@ llama_model_cohere2moe::graph::graph(const llama_model & model, const llm_graph_
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 llama_model_cohere2moe::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
@@ -426,6 +424,4 @@ llama_model_cohere2moe::graph_mtp::graph_mtp(const llama_model & model, const ll
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/command-r.cpp

@@ -139,6 +139,4 @@ llama_model_command_r::graph::graph(const llama_model & model, const llm_graph_p
     }
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/dbrx.cpp

@@ -149,6 +149,4 @@ llama_model_dbrx::graph::graph(const llama_model & model, const llm_graph_params
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/deci.cpp

@@ -186,6 +186,4 @@ llama_model_deci::graph::graph(const llama_model & model, const llm_graph_params
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/deepseek.cpp

@@ -189,6 +189,4 @@ llama_model_deepseek::graph::graph(const llama_model & model, const llm_graph_pa
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/deepseek2.cpp

@@ -406,7 +406,6 @@ llama_model_deepseek2::graph_mtp::graph_mtp(const llama_model & model, const llm
     cb(cur, "result_output", -1);
 
     res->t_logits = cur;
-    ggml_build_forward_expand(gf, cur);
 }
 
 llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_params & params) :
@@ -703,6 +702,4 @@ llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_p
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/deepseek32.cpp

@@ -475,8 +475,6 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 // LLM_GRAPH_TYPE_DECODER_MTP draft head for DeepSeek V3.2 (DEEPSEEK32).
@@ -716,6 +714,5 @@ llama_model_deepseek32::graph_mtp::graph_mtp(const llama_model & model, const ll
     cb(cur, "result_output", -1);
 
     res->t_logits = cur;
-    ggml_build_forward_expand(gf, cur);
 }
 

src/models/deepseek4.cpp

@@ -1353,8 +1353,6 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p
     cur = ggml_mul_mat(ctx0, model.output, cur);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 
@@ -1366,7 +1364,6 @@ llama_model_deepseek4::graph_mtp::graph_mtp(const llama_model & model, const llm
             cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&
             "nextn_layer_offset out of range [0, n_layer_nextn)");
     ASSERT_EMBD_OR_TOKEN(ubatch);
-    GGML_ASSERT(ubatch.token && "DEEPSEEK4 MTP requires token input");
 
     const int64_t hc = hparams.dsv4_hc_mult;
     GGML_ASSERT(hparams.n_embd_out() == (uint32_t) (n_embd*hc) && "DEEPSEEK4 MTP hidden width mismatch");
@@ -1503,5 +1500,4 @@ llama_model_deepseek4::graph_mtp::graph_mtp(const llama_model & model, const llm
     cb(cur, "result_output", -1);
 
     res->t_logits = cur;
-    ggml_build_forward_expand(gf, cur);
 }

src/models/delta-net-base.cpp

@@ -205,9 +205,6 @@ std::pair<ggml_tensor *, ggml_tensor *> llm_build_delta_net_base::build_delta_ne
             ggml_row_size(g_cs->type, g_cs->ne[0] - 1));
     cb(g_last, "g_last", il);
 
-    // TODO: remove this cont when CUDA supports non-cont unary ops
-    g_last = ggml_cont(ctx0, g_last);
-
     // [1, 1, n_chunks, H_v * n_seqs] KDA: [S_k, 1, n_chunks, H_v * n_seqs]
     ggml_tensor * g_last_exp_t = ggml_transpose(ctx0, ggml_exp(ctx0, g_last));
     cb(g_last_exp_t, "g_last_exp_t", il);

src/models/dflash.cpp

@@ -337,7 +337,7 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model &
         // bonus anchor slot: pass the logits through unbiased, pad the (unread) confidence column
         cat = ggml_cont(ctx0, ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, 0));
         if (has_conf) {
-            cat_conf = ggml_sigmoid(ctx0, ggml_cont(ctx0, ggml_view_2d(ctx0, base, 1, n_blocks, base_stride, 0)));
+            cat_conf = ggml_sigmoid(ctx0, ggml_view_2d(ctx0, base, 1, n_blocks, base_stride, 0));
         }
     }
 
@@ -841,8 +841,6 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
     cb(cur, "result_output", -1);
     res->t_logits = cur;
 
-    ggml_build_forward_expand(gf, cur);
-
     // DSpark: bias the draft logits with the Markov head
     if (model.dspark_markov_w1) {
         build_dspark_markov_head(*this, model, inp_tokens);
@@ -1022,8 +1020,6 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_
     cb(cur, "result_output", -1);
     res->t_logits = cur;
 
-    ggml_build_forward_expand(gf, cur);
-
     if (model.dspark_markov_w1) {
         build_dspark_markov_head(*this, model, inp_tokens);
     }

src/models/dots1.cpp

@@ -188,6 +188,4 @@ llama_model_dots1::graph::graph(const llama_model & model, const llm_graph_param
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/dots3note.cpp

@@ -471,6 +471,4 @@ llama_model_dots3note::graph::graph(const llama_model & model, const llm_graph_p
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/dream.cpp

@@ -133,6 +133,4 @@ llama_model_dream::graph::graph(const llama_model & model, const llm_graph_param
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/eagle3.cpp

@@ -321,8 +321,6 @@ llama_model_eagle3::graph<false>::graph(const llama_model & model, const llm_gra
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 std::unique_ptr<llm_graph_context> llama_model_eagle3::build_arch_graph(const llm_graph_params & params) const {

src/models/ernie4-5-moe.cpp

@@ -128,6 +128,4 @@ llama_model_ernie4_5_moe::graph::graph(const llama_model & model, const llm_grap
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/ernie4-5.cpp

@@ -159,6 +159,4 @@ llama_model_ernie4_5::graph::graph(const llama_model & model, const llm_graph_pa
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/exaone-moe.cpp

@@ -234,6 +234,4 @@ llama_model_exaone_moe::graph::graph(const llama_model & model, const llm_graph_
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/exaone.cpp

@@ -131,6 +131,4 @@ llama_model_exaone::graph::graph(const llama_model & model, const llm_graph_para
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/exaone4.cpp

@@ -181,8 +181,6 @@ llama_model_exaone4::graph<iswa>::graph(const llama_model & model, const llm_gra
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 // Explicit template instantiations

src/models/falcon-h1.cpp

@@ -204,6 +204,4 @@ llama_model_falcon_h1::graph::graph(const llama_model & model, const llm_graph_p
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/falcon.cpp

@@ -156,6 +156,4 @@ llama_model_falcon::graph::graph(const llama_model & model, const llm_graph_para
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/gemma.cpp

@@ -134,6 +134,4 @@ llama_model_gemma::graph::graph(const llama_model & model, const llm_graph_param
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/gemma2.cpp

@@ -170,6 +170,4 @@ llama_model_gemma2::graph::graph(const llama_model & model, const llm_graph_para
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }