llama-bench : fix OOB access of hf_file (#29515)

Signed-off-by: Adrien Gallouët <angt@huggingface.co>

.devops/musa.Dockerfile

@@ -1,10 +1,9 @@
 ARG UBUNTU_VERSION=22.04
 # This needs to generally match the container host's environment.
-ARG MUSA_VERSION=rc4.3.0
 # Target the MUSA build image
-ARG BASE_MUSA_DEV_CONTAINER=docker.io/mthreads/musa:${MUSA_VERSION}-devel-ubuntu${UBUNTU_VERSION}-amd64
+ARG BASE_MUSA_DEV_CONTAINER=registry.mthreads.com/mcconline/inference/pytorch:2.9.1.post1-py3.10-musa5.2.0-mp31-devel-ubuntu${UBUNTU_VERSION}-amd64
 
-ARG BASE_MUSA_RUN_CONTAINER=docker.io/mthreads/musa:${MUSA_VERSION}-runtime-ubuntu${UBUNTU_VERSION}-amd64
+ARG BASE_MUSA_RUN_CONTAINER=${BASE_MUSA_DEV_CONTAINER}
 
 ARG BUILD_DATE=N/A
 ARG APP_VERSION=N/A

docs/docker.md

@@ -123,7 +123,7 @@ You may want to pass in some different `ARGS`, depending on the MUSA environment
 
 The defaults are:
 
-- `MUSA_VERSION` set to `rc4.3.0`
+- the base image is the MUSA 5.2.0 image from the Moore Threads registry
 
 The resulting images, are essentially the same as the non-MUSA images:
 

docs/ops.md

@@ -14,16 +14,16 @@ Legend:
 
 | Operation | BLAS | CANN | CPU | CUDA | ET | HTP | MTL | OpenCL | SYCL | Vulkan | WebGPU | ZenDNN | zDNN |
 |-----------|------|------|------|------|------|------|------|------|------|------|------|------|------|
-|                              ABS | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              ABS | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                              ACC | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              ADD | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             ADD1 | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                           ADD_ID | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                           ARANGE | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                           ARGMAX | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                          ARGSORT | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                           ARGMAX | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                          ARGSORT | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             CEIL | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                            CLAMP | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                            CLAMP | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        COL2IM_1D | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                           CONCAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                             CONT | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
@@ -55,7 +55,7 @@ Legend:
 |                GATED_LINEAR_ATTN | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                            GEGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        GEGLU_ERF | ❌ | ✅ | ✅ | ✅ | 🟡 | ❌ | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                      GEGLU_QUICK | ❌ | ✅ | ✅ | ✅ | 🟡 | ❌ | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                      GEGLU_QUICK | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             GELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                         GELU_ERF | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                       GELU_QUICK | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
@@ -64,17 +64,21 @@ Legend:
 |                       GROUP_NORM | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                      HARDSIGMOID | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        HARDSWISH | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                           IM2COL | ❌ | ✅ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                           IM2COL | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        IM2COL_3D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                          L2_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | 🟡 | ❌ | ❌ |
-|                       LEAKY_RELU | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | 🟡 | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                       LEAKY_RELU | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | 🟡 | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                LIGHTNING_INDEXER | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                              LOG | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              LOG | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             MEAN | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ | ❌ |
 |                              MUL | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          MUL_MAT | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 |
-|                 MUL_MAT_HADAMARD | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                 MUL_MAT_HADAMARD | ❌ | ❌ | ❌ | ❌ | ✅ | 🟡 | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                       MUL_MAT_ID | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | 🟡 | ❌ |
+|                  MUL_MAT_ID_W4A4 | ❌ | ❌ | ❌ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                  MUL_MAT_ID_W4A8 | ❌ | ❌ | ❌ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                     MUL_MAT_W4A4 | ❌ | ❌ | ❌ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                     MUL_MAT_W4A8 | ❌ | ❌ | ❌ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                              NEG | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                   OPT_STEP_ADAMW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -85,12 +89,12 @@ Legend:
 |                          POOL_1D | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                          POOL_2D | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                            REGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | ❌ | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             RELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             RELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                           REPEAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                      REPEAT_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                         RMS_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                    RMS_NORM_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                             ROLL | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                             ROLL | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                             ROPE | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        ROPE_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                            ROUND | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
@@ -108,20 +112,20 @@ Legend:
 |                         SOFT_MAX | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                    SOFT_MAX_BACK | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ | ❌ | 🟡 | ✅ | ❌ | ❌ | ❌ |
 |                        SOLVE_TRI | ❌ | ❌ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ |
-|                              SQR | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             SQRT | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              SQR | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             SQRT | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                         SSM_CONV | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                         SSM_SCAN | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | 🟡 | 🟡 | ✅ | ❌ | ❌ |
-|                             STEP | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             STEP | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                              SUB | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                              SUM | ❌ | 🟡 | ✅ | 🟡 | ❌ | ❌ | 🟡 | ❌ | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
+|                              SUM | ❌ | 🟡 | ✅ | 🟡 | ❌ | 🟡 | 🟡 | ❌ | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
 |                         SUM_ROWS | ❌ | ✅ | ✅ | 🟡 | ❌ | 🟡 | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ |
 |                           SWIGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                     SWIGLU_CLAMP | ❌ | ❌ | ❌ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                       SWIGLU_OAI | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             TANH | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |               TIMESTEP_EMBEDDING | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                            TOP_K | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ✅ | ❌ | ✅ | 🟡 | ✅ | ❌ | ❌ |
+|                            TOP_K | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | 🟡 | ✅ | ❌ | ❌ |
 |                              TRI | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                            TRUNC | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          UPSCALE | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |

ggml/src/ggml-cuda/fwht.cu

@@ -1,9 +1,10 @@
 #include "common.cuh"
+#include "convert.cuh"
 #include "fwht.cuh"
 
-template <int N>
+template <int N, typename T>
 __launch_bounds__(4*ggml_cuda_get_physical_warp_size(), 1)
-__global__ void fwht_cuda(const float * src, float * dst, const int64_t n_rows, const float scale) {
+__global__ void fwht_cuda(const T * src, float * dst, const int64_t n_rows, const float scale) {
     constexpr int warp_size = ggml_cuda_get_physical_warp_size();
 
     const int64_t r = (int64_t) blockIdx.x * blockDim.y + threadIdx.y;
@@ -22,7 +23,7 @@ __global__ void fwht_cuda(const float * src, float * dst, const int64_t n_rows,
     ggml_cuda_pdl_sync();
 #pragma unroll
     for (int i = 0; i < el_w; ++i) {
-        reg[i] = src[i * warp_size + lane] * scale;
+        reg[i] = ggml_cuda_cast<float>(src[i * warp_size + lane]) * scale;
     }
 
 #pragma unroll
@@ -58,15 +59,12 @@ __global__ void fwht_cuda(const float * src, float * dst, const int64_t n_rows,
     }
 }
 
-bool ggml_cuda_op_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst) {
-    GGML_ASSERT(ggml_are_same_shape(src, dst));
-    if (!ggml_is_contiguous(src) || !ggml_is_contiguous(dst)) {
-        return false;
-    }
+template <typename T>
+static bool ggml_cuda_op_fwht_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst) {
     const int     n    = src->ne[0];
     const int64_t rows = ggml_nrows(src);
 
-    const float * src_d = (const float *) src->data;
+    const T *     src_d = (const T *) src->data;
     float *       dst_d = (float *) dst->data;
 
     const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
@@ -84,18 +82,47 @@ bool ggml_cuda_op_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * src,
 
     switch (n) {
         case 64:
-            ggml_cuda_kernel_launch(fwht_cuda<64>, launch_params, src_d, dst_d, rows, scale);
+            ggml_cuda_kernel_launch(fwht_cuda<64, T>, launch_params, src_d, dst_d, rows, scale);
             return true;
         case 128:
-            ggml_cuda_kernel_launch(fwht_cuda<128>, launch_params, src_d, dst_d, rows, scale);
+            ggml_cuda_kernel_launch(fwht_cuda<128, T>, launch_params, src_d, dst_d, rows, scale);
             return true;
         case 256:
-            ggml_cuda_kernel_launch(fwht_cuda<256>, launch_params, src_d, dst_d, rows, scale);
+            ggml_cuda_kernel_launch(fwht_cuda<256, T>, launch_params, src_d, dst_d, rows, scale);
             return true;
         case 512:
-            ggml_cuda_kernel_launch(fwht_cuda<512>, launch_params, src_d, dst_d, rows, scale);
+            ggml_cuda_kernel_launch(fwht_cuda<512, T>, launch_params, src_d, dst_d, rows, scale);
             return true;
         default:
             return false;
     }
 }
+
+bool ggml_cuda_op_mul_mat_use_fwht(const struct ggml_tensor * op) {
+    const struct ggml_tensor * a = op->src[0];
+    const struct ggml_tensor * b = op->src[1];
+
+    return op->op == GGML_OP_MUL_MAT && ggml_get_op_params_i32(op, 1) == GGML_HINT_SRC0_IS_HADAMARD &&
+           a->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 &&
+           (b->type == GGML_TYPE_F32 || b->type == GGML_TYPE_F16) && ggml_is_contiguous(b) && ggml_is_contiguous(op) &&
+           ggml_are_same_shape(b, op);
+}
+
+bool ggml_cuda_op_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst) {
+    GGML_ASSERT(ggml_are_same_shape(src, dst));
+    if (!ggml_is_contiguous(src) || !ggml_is_contiguous(dst)) {
+        return false;
+    }
+    if (dst->type != GGML_TYPE_F32) {
+        return false;
+    }
+
+    switch (src->type) {
+        case GGML_TYPE_F32:
+            return ggml_cuda_op_fwht_impl<float>(ctx, src, dst);
+        case GGML_TYPE_F16:
+            return ggml_cuda_op_fwht_impl<half>(ctx, src, dst);
+        default:
+            return false;
+    }
+}

ggml/src/ggml-cuda/fwht.cuh

@@ -1,4 +1,6 @@
 #include "common.cuh"
 
+bool ggml_cuda_op_mul_mat_use_fwht(const struct ggml_tensor * op);
+
 // Returns whether the Fast Walsh-Hadamard transform could be used.
 bool ggml_cuda_op_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst);

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

@@ -1818,8 +1818,7 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) {
 static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
     GGML_TENSOR_BINARY_OP_LOCALS
 
-    const int32_t hint = ggml_get_op_params_i32(dst, 1);
-    if (hint == GGML_HINT_SRC0_IS_HADAMARD && ggml_cuda_op_fwht(ctx, src1, dst)) {
+    if (ggml_cuda_op_mul_mat_use_fwht(dst) && ggml_cuda_op_fwht(ctx, src1, dst)) {
         return;
     }
 
@@ -5212,7 +5211,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                 if (a->nb[0] != ggml_element_size(a) || b->nb[0] != ggml_element_size(b)) {
                     return false; // TODO this could in principle be implemented though currently there is no use case.
                 }
-                if (b->type == GGML_TYPE_F16 && a->type != GGML_TYPE_F16) {
+                if (b->type == GGML_TYPE_F16 && a->type != GGML_TYPE_F16 && !ggml_cuda_op_mul_mat_use_fwht(op)) {
                     return false;
                 }
                 if (op->op == GGML_OP_MUL_MAT_ID && ggml_get_op_params_i32(op, 3) == GGML_PREC_F32) {
@@ -5488,8 +5487,9 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
 
             if (op->src[3]->ne[0] == 1) {
                 // Mamba2
-                // (kernel only supports (d_state == 128 || d_state == 256) && d_head % 16 == 0)
-                return (op->src[0]->ne[0] == 128 || op->src[0]->ne[0] == 256) && op->src[0]->ne[1] % 16 == 0;
+                // (kernel only supports (d_state == 96 || d_state == 128 || d_state == 256) && d_head % 16 == 0)
+                const int64_t d_state = op->src[0]->ne[0];
+                return (d_state == 96 || d_state == 128 || d_state == 256) && op->src[0]->ne[1] % 16 == 0;
             } else {
                 if (K > 1) {
                     return false;

ggml/src/ggml-cuda/ssm-scan.cu

@@ -163,13 +163,19 @@ __global__ void __launch_bounds__(d_state, 1)
     const int lane     = threadIdx.x % WARP_SIZE;
     const int warp_idx = blockIdx.x  * c_factor + warp;
 
+    ggml_cuda_pdl_sync();
+
+    // the last block can have unused warps when n_head*d_head is not a multiple of c_factor
+    if (warp_idx >= n_head * d_head) {
+        return;
+    }
+
     const int head_idx =  warp_idx / d_head;
     const int head_off = (warp_idx % d_head) * sizeof(float);
     const int seq_idx  = blockIdx.y;
 
     const int group_off = (head_idx / (n_head / n_group)) * d_state * sizeof(float);
 
-    ggml_cuda_pdl_sync();
     // TODO: refactor strides to be in elements/floats instead of bytes to be cleaner and consistent with the rest of the codebase
     const float * s0_warp = (const float *) ((const char *) src0 + src6[seq_idx] * src0_nb3 + head_idx * src0_nb2 + head_off * d_state);
     const float * x_warp  = (const float *) ((const char *) src1 + (seq_idx * src1_nb3) + (warp_idx * sizeof(float)));
@@ -246,7 +252,17 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
     // NOTE: if you change conditions here, be sure to update the corresponding supports_op condition!
     if (src3_nb1 == sizeof(float)) {
         // Mamba-2
-        if (d_state == 128) {
+        if (d_state == 96) {
+            constexpr int threads   = 96;
+            constexpr int num_warps = threads/WARP_SIZE;
+
+            const dim3 blocks((n_head * head_dim + (num_warps - 1)) / num_warps, n_seq, 1);
+            const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks, threads, 0, stream);
+            ggml_cuda_kernel_launch(ssm_scan_f32_group<96/WARP_SIZE, 96>, launch_params,
+                    src0, src1, src2, src3, src4, src5, src6, dst,
+                    src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
+                    src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K);
+        } else if (d_state == 128) {
             constexpr int threads   = 128;
             constexpr int num_warps = threads/WARP_SIZE;
 
@@ -267,7 +283,7 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
                     src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1,
                     src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K);
         } else {
-            GGML_ABORT("doesn't support d_state!=(128 or 256).");
+            GGML_ABORT("doesn't support d_state!=(96, 128 or 256).");
         }
     } else {
         // Mamba-1

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

@@ -63,6 +63,7 @@
 #include "htp/rope-ops.h"
 #include "htp/ssm-conv.h"
 #include "htp/gated-delta-net-ops.h"
+#include "htp/argsort-ops.h"
 #include "htp_iface.h"
 #include "htp-drv.h"
 
@@ -368,6 +369,13 @@ static void ggml_hexagon_precompute_gated_delta_net_params(
     struct htp_gdn_kernel_params * kparams
 );
 
+static void ggml_hexagon_precompute_sort_params(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op,
+    bool is_top_k,
+    struct htp_sort_kernel_params * kparams
+);
+
 static void ggml_hexagon_precompute_fused_mmnx_params(
     const struct ggml_hexagon_session * sess,
     const struct ggml_tensor * src0,
@@ -4998,13 +5006,48 @@ static bool ggml_hexagon_precompute_binary_params(
     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 &&
+        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 &&
         src1->ne[0] == src0->ne[0] &&
-        (src1->ne[1] == src0->ne[1] || src1->ne[1] == 1) &&
+        src1->ne[1] == src0->ne[1] &&
         (src1->ne[2] == src0->ne[2] || src1->ne[2] == 1) &&
         (src1->ne[3] == src0->ne[3] || src1->ne[3] == 1);
-    const bool is_row_bcast = is_same_shape && src1->ne[1] == 1 && src1->ne[2] == 1 && src1->ne[3] == 1;
-    const bool is_complex   = !is_add_id && !is_scalar && !is_same_shape && (src1->ne[0] == src0->ne[0]);
+    const bool is_complex   = !is_add_id && !is_scalar && !is_same_shape && !is_row_bcast && (src1->ne[0] == src0->ne[0]);
+    const bool is_contig    = ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(dst);
+    const bool is_scalar_broadcast = !is_add_id && (ggml_nelements(src1) == 1);
+
+    if (!is_add_id && is_contig && (ggml_are_same_shape(src0, src1) || is_scalar_broadcast)) {
+        const uint32_t total_elems = (uint32_t) ggml_nelements(src0);
+        const uint32_t n_threads = sess->n_threads;
+        const uint32_t max_chunk_elems = 32768 / elem_size;
+        const uint32_t min_chunk_elems = 256;
+        const uint32_t target_chunk_elems = hex_round_up((total_elems + (2 * n_threads) - 1) / (2 * n_threads), 32);
+        const uint32_t chunk_size = (std::min)(max_chunk_elems, (std::max)(target_chunk_elems, min_chunk_elems));
+        const uint32_t chunk_bytes = hex_round_up(chunk_size * elem_size, 128);
+
+        kparams->kernel_type           = HTP_BINARY_KERNEL_CHUNKED;
+        kparams->n_threads             = n_threads;
+        kparams->rows_per_buffer       = 1;
+        kparams->src0_row_size_aligned = src0_row_size_aligned;
+        kparams->src1_row_size_aligned = is_scalar_broadcast ? 0 : src1_row_size_aligned;
+        kparams->dst_row_size_aligned  = dst_row_size_aligned;
+        kparams->src1_size             = 0;
+        kparams->chunk_size            = chunk_size;
+        kparams->chunk_bytes           = chunk_bytes;
+        kparams->is_scalar             = is_scalar_broadcast ? 1 : 0;
+
+        struct htp_binary_vtcm_layout L;
+        htp_binary_vtcm_layout_build(&L, kparams, sess->vtcm_size);
+        if (L.total_bytes == 0 || L.total_bytes > sess->vtcm_size) {
+            return false;
+        }
+
+        kparams->vtcm_size = L.total_bytes;
+        return true;
+    }
 
     enum htp_binary_kernel_type kernel_type;
     size_t src1_size = 0;
@@ -5042,6 +5085,32 @@ static bool ggml_hexagon_precompute_binary_params(
     struct htp_binary_vtcm_layout L;
     htp_binary_vtcm_layout_build(&L, kparams, sess->vtcm_size);
     if (L.rows_per_buffer == 0 || L.total_bytes > sess->vtcm_size) {
+        if (!is_add_id && is_contig && (ggml_are_same_shape(src0, src1) || is_scalar_broadcast)) {
+            const uint32_t total_elems = (uint32_t) ggml_nelements(src0);
+            const uint32_t n_threads = sess->n_threads;
+            const uint32_t max_chunk_elems = 32768 / elem_size;
+            const uint32_t min_chunk_elems = 256;
+            const uint32_t target_chunk_elems = hex_round_up((total_elems + (2 * n_threads) - 1) / (2 * n_threads), 32);
+            const uint32_t chunk_size = (std::min)(max_chunk_elems, (std::max)(target_chunk_elems, min_chunk_elems));
+            const uint32_t chunk_bytes = hex_round_up(chunk_size * elem_size, 128);
+
+            kparams->kernel_type           = HTP_BINARY_KERNEL_CHUNKED;
+            kparams->n_threads             = n_threads;
+            kparams->rows_per_buffer       = 1;
+            kparams->src1_row_size_aligned = is_scalar_broadcast ? 0 : src1_row_size_aligned;
+            kparams->src1_size             = 0;
+            kparams->chunk_size            = chunk_size;
+            kparams->chunk_bytes           = chunk_bytes;
+            kparams->is_scalar             = is_scalar_broadcast ? 1 : 0;
+
+            htp_binary_vtcm_layout_build(&L, kparams, sess->vtcm_size);
+            if (L.total_bytes == 0 || L.total_bytes > sess->vtcm_size) {
+                return false;
+            }
+
+            kparams->vtcm_size = L.total_bytes;
+            return true;
+        }
         return false;
     }
 
@@ -5139,49 +5208,51 @@ static void ggml_hexagon_precompute_get_rows_params(
     const uint32_t ne12 = src1->ne[2];
     const uint32_t nr = ne10 * ne11 * ne12;
 
-    const size_t nb01 = src0->nb[1];
-    const size_t nb1 = dst->nb[1];
+    const ggml_tensor * src0_base = src0->view_src ? src0->view_src : src0;
+    const auto * extra = src0_base->buffer && ggml_backend_buffer_is_hexagon(src0_base->buffer) ?
+        (const ggml_hexagon_tensor_extra *) src0_base->extra : nullptr;
+    const bool tiled = src0->type == GGML_TYPE_Q4_0 || (extra && (extra->flags & GGML_HEXAGON_TENSOR_REPACK) != 0) ||
+                       sess->needs_repack.count(src0_base) || sess->needs_repack.count(src0);
 
-    const bool can_use_dma = (src0->type == dst->type) && (nb01 == nb1);
-    const bool use_dma = can_use_dma && (ne00 >= 2048);
-
-    kparams->use_dma = use_dma ? 1 : 0;
-
-    uint32_t chunks_per_row = 1;
-    uint32_t chunk_size = ne00;
-    uint32_t total_tasks = nr;
-
-    if (use_dma) {
-        kparams->n_threads = (std::min)((uint32_t)sess->n_threads, nr);
-        kparams->tasks_per_thread = (nr + kparams->n_threads - 1) / kparams->n_threads;
+    if (src0->type == dst->type) {
+        kparams->kernel_type = HTP_GET_ROWS_KERNEL_SAMETYPE;
+    } else if (tiled) {
+        kparams->kernel_type = HTP_GET_ROWS_KERNEL_TILED;
     } else {
-        if (src0->type == GGML_TYPE_F32 && nr < sess->n_threads) {
-            const uint32_t min_chunk_size = 1024;
-            uint32_t max_chunks = ne00 / min_chunk_size;
-            if (max_chunks == 0) {
-                max_chunks = 1;
-            }
-            chunks_per_row = (std::min)((sess->n_threads + nr - 1) / nr, max_chunks);
-            chunk_size = (ne00 + chunks_per_row - 1) / chunks_per_row;
-            total_tasks = nr * chunks_per_row;
-        }
-        kparams->n_threads = (std::min)(total_tasks, (uint32_t)sess->n_threads);
-        kparams->tasks_per_thread = (total_tasks + kparams->n_threads - 1) / kparams->n_threads;
+        kparams->kernel_type = HTP_GET_ROWS_KERNEL_FLAT;
     }
 
+    const uint32_t chunks_per_row = 1;
+    const uint32_t chunk_size = ne00;
+    const uint32_t total_tasks = nr;
+
+    kparams->n_threads = (std::min)((uint32_t)sess->n_threads, total_tasks);
+
+    struct htp_get_rows_vtcm_layout vtcm_layout = {};
+    while (kparams->n_threads > 0) {
+        htp_get_rows_vtcm_layout_build(&vtcm_layout, kparams->kernel_type, src0->type, ne00, kparams->n_threads);
+        if (vtcm_layout.total_bytes <= sess->vtcm_size) {
+            break;
+        }
+        --kparams->n_threads;
+    }
+
+    if (kparams->n_threads == 0 && total_tasks > 0) {
+        htp_get_rows_vtcm_layout_build(&vtcm_layout, kparams->kernel_type, src0->type, ne00, 1);
+    }
+
+    kparams->vtcm_size = (total_tasks == 0) ? 0 : vtcm_layout.total_bytes;
+    kparams->tasks_per_thread = kparams->n_threads > 0 ? (total_tasks + kparams->n_threads - 1) / kparams->n_threads : 0;
+
     kparams->chunks_per_row = chunks_per_row;
     kparams->chunk_size = chunk_size;
     kparams->total_tasks = total_tasks;
 
-    kparams->div_ne10 = init_fastdiv_values(ne10);
-    kparams->div_ne10_ne11 = init_fastdiv_values(ne10 * ne11);
-    kparams->div_chunks_per_row = init_fastdiv_values(chunks_per_row);
-    kparams->div_ne02 = init_fastdiv_values(ne02);
-    kparams->div_ne03 = init_fastdiv_values(ne03);
-
-    struct htp_get_rows_vtcm_layout vtcm_layout;
-    htp_get_rows_vtcm_layout_build(&vtcm_layout, src0->type, ne00, kparams->n_threads);
-    kparams->vtcm_size = vtcm_layout.total_bytes;
+    kparams->div_ne10 = ne10 > 0 ? init_fastdiv_values(ne10) : fastdiv_values{0, 0};
+    kparams->div_ne10_ne11 = (ne10 * ne11) > 0 ? init_fastdiv_values(ne10 * ne11) : fastdiv_values{0, 0};
+    kparams->div_chunks_per_row = chunks_per_row > 0 ? init_fastdiv_values(chunks_per_row) : fastdiv_values{0, 0};
+    kparams->div_ne02 = ne02 > 0 ? init_fastdiv_values(ne02) : fastdiv_values{0, 0};
+    kparams->div_ne03 = ne03 > 0 ? init_fastdiv_values(ne03) : fastdiv_values{0, 0};
 }
 
 static void ggml_hexagon_precompute_set_rows_params(
@@ -5479,6 +5550,52 @@ static void ggml_hexagon_precompute_gated_delta_net_params(
     if (kparams->n_threads > 0) kparams->div_n_threads = init_fastdiv_values(kparams->n_threads);
 }
 
+static void ggml_hexagon_precompute_sort_params(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op,
+    bool is_top_k,
+    struct htp_sort_kernel_params * kparams
+) {
+    memset(kparams, 0, sizeof(*kparams));
+
+    const struct ggml_tensor * src0 = op->src[0];
+    const struct ggml_tensor * dst  = op;
+
+    const uint32_t total_rows = src0->ne[1] * src0->ne[2] * src0->ne[3];
+    const uint32_t ne00       = src0->ne[0];
+    const uint32_t k          = dst->ne[0];
+
+    int32_t order = GGML_SORT_ORDER_DESC;
+    if (!is_top_k) {
+        order = ((const int32_t *) op->op_params)[0];
+    }
+
+    const uint32_t n_threads_max = sess->n_threads > 0 ? sess->n_threads : 4;
+    const size_t vtcm_budget = sess->vtcm_size > 0 ? sess->vtcm_size : (8 * 1024 * 1024);
+
+    struct htp_sort_vtcm_layout layout;
+    bool ok = htp_sort_solve_layout(&layout, ne00, total_rows, k, n_threads_max, vtcm_budget, is_top_k);
+    GGML_ASSERT(ok);
+
+    kparams->n_threads         = (int32_t) layout.n_threads;
+    kparams->total_rows        = (int32_t) total_rows;
+    kparams->row_start         = 0;
+    kparams->row_end           = (int32_t) total_rows;
+    kparams->ne00              = (int32_t) ne00;
+    kparams->k                 = (int32_t) k;
+    kparams->order             = order;
+    kparams->is_top_k          = is_top_k ? 1 : 0;
+    kparams->use_dma           = 1;
+    kparams->chunk_elems       = (int32_t) layout.chunk_elems;
+    kparams->n_chunks          = (int32_t) layout.n_chunks;
+    kparams->vtcm_size         = (int32_t) layout.total_bytes;
+    kparams->phase1_slot_size  = (int32_t) layout.phase1_slot_size;
+    kparams->merge_values_off  = (int32_t) layout.merge_values_off;
+    kparams->merge_indices_off = (int32_t) layout.merge_indices_off;
+    kparams->merge_elems       = (int32_t) layout.merge_elems;
+    kparams->n_slots           = (int32_t) layout.n_slots;
+}
+
 static void ggml_hexagon_precompute_fused_mmnx_params(
     const struct ggml_hexagon_session * sess,
     const struct ggml_tensor * src0, // W0
@@ -5830,7 +5947,8 @@ static bool ggml_hexagon_supported_unary(const struct ggml_hexagon_session * ses
             case GGML_OP_LOG:
                 break;
             case GGML_OP_UNARY:
-                if (ggml_get_unary_op(op) != GGML_UNARY_OP_ABS) {
+                if (ggml_get_unary_op(op) != GGML_UNARY_OP_ABS &&
+                    ggml_get_unary_op(op) != GGML_UNARY_OP_STEP) {
                     return false;
                 }
                 break;
@@ -5853,6 +5971,26 @@ static bool ggml_hexagon_supported_unary(const struct ggml_hexagon_session * ses
     GGML_UNUSED(sess);
 }
 
+static bool ggml_hexagon_supported_sum(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
+    const struct ggml_tensor * src0 = op->src[0];
+    const struct ggml_tensor * dst  = op;
+
+    if (src0->type != GGML_TYPE_F32) {
+        return false;
+    }
+    if (dst->type != GGML_TYPE_F32) {
+        return false;
+    }
+
+    if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(dst)) {
+        return false;
+    }
+
+    return true;
+
+    GGML_UNUSED(sess);
+}
+
 static bool ggml_hexagon_supported_sum_rows(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
     const struct ggml_tensor * src0 = op->src[0];
     const struct ggml_tensor * dst  = op;
@@ -5874,6 +6012,26 @@ static bool ggml_hexagon_supported_sum_rows(const struct ggml_hexagon_session *
     GGML_UNUSED(sess);
 }
 
+static bool ggml_hexagon_supported_argmax(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
+    const struct ggml_tensor * src0 = op->src[0];
+    const struct ggml_tensor * dst  = op;
+
+    if (src0->type != GGML_TYPE_F32) {
+        return false;
+    }
+    if (dst->type != GGML_TYPE_I32) {
+        return false;
+    }
+
+    if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(dst)) {
+        return false;
+    }
+
+    return true;
+
+    GGML_UNUSED(sess);
+}
+
 static bool ggml_hexagon_supported_activations(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
     const struct ggml_tensor * src0 = op->src[0];
     const struct ggml_tensor * src1 = op->src[1];
@@ -6020,19 +6178,36 @@ static bool ggml_hexagon_supported_get_rows(const struct ggml_hexagon_session *
     const struct ggml_tensor * src1 = op->src[1]; // indices
     const struct ggml_tensor * dst  = op;
 
-    if (src0->extra) {
-        const auto * extra = (const ggml_hexagon_tensor_extra *) src0->extra;
-        if (extra->flags & GGML_HEXAGON_TENSOR_REPACK) {
+    if (src0->type == GGML_TYPE_Q4_0 && src0->view_src) {
+        return false;
+    }
+
+    const ggml_tensor * src0_base = src0->view_src ? src0->view_src : src0;
+    bool is_repacked = false;
+    if (src0_base->buffer && ggml_backend_buffer_is_hexagon(src0_base->buffer) && src0_base->extra) {
+        const auto * extra = (const ggml_hexagon_tensor_extra *) src0_base->extra;
+        is_repacked = (extra->flags & GGML_HEXAGON_TENSOR_REPACK) != 0;
+        if (is_repacked && src0->type != GGML_TYPE_Q4_0 && src0->type != GGML_TYPE_Q8_0) {
             return false;
         }
     }
+    is_repacked = is_repacked || sess->needs_repack.count(src0_base) || sess->needs_repack.count(src0);
 
-    if (src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_I32 && src0->ne[0] < 32) {
+    // View offsets use the raw quantized layout and cannot address a tiled allocation.
+    if (src0->view_src && is_repacked) {
+        return false;
+    }
+
+    if (src0->type == GGML_TYPE_Q4_0 && src0->buffer && !is_repacked) {
+        return false;
+    }
+
+    if (src0->type != dst->type && src0->ne[0] < 32) {
         return false;
     }
 
     if (src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_F16 &&
-        src0->type != GGML_TYPE_Q8_0 && src0->type != GGML_TYPE_I32) {
+        src0->type != GGML_TYPE_Q4_0 && src0->type != GGML_TYPE_Q8_0 && src0->type != GGML_TYPE_I32) {
         return false;
     }
 
@@ -6040,13 +6215,30 @@ static bool ggml_hexagon_supported_get_rows(const struct ggml_hexagon_session *
         return false;
     }
 
-    if (src0->type == GGML_TYPE_I32) {
-        if (dst->type != GGML_TYPE_I32) {
+    if (src0->type == dst->type) {
+        if (src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_I32 && src0->type != GGML_TYPE_F16) {
             return false;
         }
-    }
-    else if (dst->type != GGML_TYPE_F32) {
+    } else if (src0->type == GGML_TYPE_I32) {
         return false;
+    } else if (dst->type != GGML_TYPE_F32) {
+        return false;
+    }
+
+    // Empty recurrent-state gathers are skipped at execution; do not split the graph for them.
+    if (ggml_is_empty(op)) {
+        return true;
+    }
+
+    struct htp_get_rows_kernel_params kparams;
+    ggml_hexagon_precompute_get_rows_params(sess, src0, src1, dst, &kparams);
+    if (kparams.n_threads == 0 || (size_t) kparams.vtcm_size > sess->vtcm_size) {
+        return false;
+    }
+
+    // Q4_0 has no raw fallback. Mark only accepted tensors for repacking.
+    if (src0->type == GGML_TYPE_Q4_0 && !src0->buffer) {
+        sess->needs_repack.insert(src0);
     }
 
     return true;
@@ -6058,42 +6250,36 @@ static bool ggml_hexagon_supported_argsort(const struct ggml_hexagon_session * s
     const struct ggml_tensor * src0 = op->src[0]; // values
     const struct ggml_tensor * dst  = op;         // indices
 
-    if (src0->type != GGML_TYPE_F32) {
+    if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_I32) {
         return false;
     }
 
-    if (dst->type != GGML_TYPE_I32) {
-        return false;
-    }
+    const uint32_t total_rows = src0->ne[1] * src0->ne[2] * src0->ne[3];
+    const uint32_t n_threads_max = sess->n_threads > 0 ? sess->n_threads : 4;
+    const size_t vtcm_budget = sess->vtcm_size > 0 ? sess->vtcm_size : (8 * 1024 * 1024);
 
-    if (src0->ne[0] > (16*1024)) {
-        // reject tensors with huge rows for now
+    struct htp_sort_vtcm_layout layout;
+    if (!htp_sort_solve_layout(&layout, src0->ne[0], total_rows, dst->ne[0], n_threads_max, vtcm_budget, false)) {
         return false;
     }
 
     return true;
-
-    GGML_UNUSED(sess);
 }
 
 static bool ggml_hexagon_supported_top_k(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) {
     const struct ggml_tensor * src0 = op->src[0]; // values
     const struct ggml_tensor * dst  = op;         // indices
 
-    if (src0->type != GGML_TYPE_F32) {
+    if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_I32) {
         return false;
     }
 
-    if (dst->type != GGML_TYPE_I32) {
-        return false;
-    }
+    const uint32_t total_rows = src0->ne[1] * src0->ne[2] * src0->ne[3];
+    const uint32_t n_threads_max = sess->n_threads > 0 ? sess->n_threads : 4;
+    const size_t vtcm_budget = sess->vtcm_size > 0 ? sess->vtcm_size : (8 * 1024 * 1024);
 
-    // Single row uses the threaded chunk+merge path. Multi-row uses one full
-    // buffer per thread, so it keeps the tighter 64K cap.
-    const bool single_row = (src0->ne[1] == 1 && src0->ne[2] == 1 && src0->ne[3] == 1);
-    const int64_t max_ne00 = single_row ? (256*1024) : (64*1024);
-
-    if (src0->ne[0] > max_ne00) {
+    struct htp_sort_vtcm_layout layout;
+    if (!htp_sort_solve_layout(&layout, src0->ne[0], total_rows, dst->ne[0], n_threads_max, vtcm_budget, true)) {
         return false;
     }
 
@@ -6401,9 +6587,11 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) {
         case GGML_OP_CONT:            return HTP_OP_CPY;
         case GGML_OP_GET_ROWS:        return HTP_OP_GET_ROWS;
         case GGML_OP_SET_ROWS:        return HTP_OP_SET_ROWS;
+        case GGML_OP_SUM:             return HTP_OP_SUM;
         case GGML_OP_SUM_ROWS:        return HTP_OP_SUM_ROWS;
         case GGML_OP_ARGSORT:         return HTP_OP_ARGSORT;
         case GGML_OP_TOP_K:           return HTP_OP_TOP_K;
+        case GGML_OP_ARGMAX:          return HTP_OP_ARGMAX;
         case GGML_OP_NORM:            return HTP_OP_NORM;
         case GGML_OP_L2_NORM:         return HTP_OP_L2_NORM;
         case GGML_OP_RMS_NORM:        return HTP_OP_RMS_NORM;
@@ -6440,6 +6628,7 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) {
                 case GGML_UNARY_OP_TANH:       return HTP_OP_UNARY_TANH;
                 case GGML_UNARY_OP_ABS:        return HTP_OP_UNARY_ABS;
                 case GGML_UNARY_OP_RELU:       return HTP_OP_UNARY_RELU;
+                case GGML_UNARY_OP_STEP:       return HTP_OP_UNARY_STEP;
             default:
                 break;
             }
@@ -6670,6 +6859,12 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg
                     node.node,
                     (struct htp_gdn_kernel_params *)node.kernel_params
                 );
+            } else if (node.opcode == HTP_OP_ARGSORT || node.opcode == HTP_OP_TOP_K) {
+                ggml_hexagon_precompute_sort_params(sess,
+                    node.node,
+                    node.opcode == HTP_OP_TOP_K,
+                    (struct htp_sort_kernel_params *) node.kernel_params
+                );
             }
             computed_nodes.push_back(std::move(node));
         }
@@ -7308,16 +7503,25 @@ static bool ggml_hexagon_supported_cpy(const struct ggml_hexagon_session * sess,
     if (dst->type != GGML_TYPE_F32 && dst->type != GGML_TYPE_F16 &&
         dst->type != GGML_TYPE_I32) return false;
 
+    const bool is_scalar  = (ggml_nelements(src0) == 1 && ggml_nelements(dst) == 1);
     const bool sametype   = (src0->type == dst->type);
-    const bool transposed = ggml_is_transposed(src0) || ggml_is_transposed(dst);
-    const bool sameshape  = !transposed && ggml_are_same_shape(src0, dst);
+    const bool transposed = !is_scalar && (ggml_is_transposed(src0) || ggml_is_transposed(dst));
+    const bool sameshape  = is_scalar || (!transposed && ggml_are_same_shape(src0, dst));
 
-    // Same-type copies also support I32.
+    if (src0->type == GGML_TYPE_I32 || dst->type == GGML_TYPE_I32) {
+        if (!sameshape) return false;
+        if (sametype) return true;
+        if ((src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_I32) ||
+            (src0->type == GGML_TYPE_I32 && dst->type == GGML_TYPE_F32)) {
+            return true;
+        }
+        return false;
+    }
+
+    // can handle any shape and any same-type (pretty slow if reshaping is required)
     if (sametype) return true;
 
-    // Type conversion is only supported between F32 and F16.
-    if (src0->type == GGML_TYPE_I32 || dst->type == GGML_TYPE_I32) return false;
-
+    // cannot handle re-shaping and type conversion at the same time
     if (!sameshape) return false;
 
     return true;
@@ -7466,10 +7670,18 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
             supp = ggml_hexagon_supported_unary(sess, op);
             break;
 
+        case GGML_OP_SUM:
+            supp = ggml_hexagon_supported_sum(sess, op);
+            break;
+
         case GGML_OP_SUM_ROWS:
             supp = ggml_hexagon_supported_sum_rows(sess, op);
             break;
 
+        case GGML_OP_ARGMAX:
+            supp = ggml_hexagon_supported_argmax(sess, op);
+            break;
+
         case GGML_OP_SOFT_MAX:
             supp = ggml_hexagon_supported_softmax(sess, op);
             break;
@@ -7486,6 +7698,7 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
                 case GGML_UNARY_OP_GELU:
                 case GGML_UNARY_OP_GELU_QUICK:
                 case GGML_UNARY_OP_RELU:
+                case GGML_UNARY_OP_STEP:
                     supp = ggml_hexagon_supported_unary(sess, op);
                     break;
                 default:

ggml/src/ggml-hexagon/htp-opnode.h

@@ -19,6 +19,8 @@
 #include "htp/ssm-conv.h"
 #include "htp/gated-delta-net-ops.h"
 #include "htp/softmax-ops.h"
+#include "htp/argsort-ops.h"
+#include "htp/get-rows-ops.h"
 
 struct htp_opnode {
     ggml_tensor * node   { nullptr };
@@ -359,17 +361,31 @@ struct htp_opformat {
         } else if (node.opcode == HTP_OP_GATED_DELTA_NET) {
             const auto * kparams = (const struct htp_gdn_kernel_params *) node.kernel_params;
             const char * path = (kparams->kernel_type == HTP_GDN_KERNEL_HMX_CHUNKED) ? "hmx-chunked" : "hvx-recurrent";
-            snprintf(str, max_size, "%s-%s vtcm %u",
-                     path,
-                     kparams->kda ? "kda" : "scalar",
+            snprintf(str, max_size, "%s-%s vtcm %u", path, kparams->kda ? "kda" : "scalar",
                      (unsigned int) (kparams->vtcm_size ? kparams->vtcm_size : kparams->vtcm_per_thread * kparams->n_threads));
         } 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 auto * kparams = (const struct htp_binary_kernel_params *) node.kernel_params;
             snprintf(str, max_size, "vtcm %u", (unsigned int) kparams->vtcm_size);
+        } else if (node.opcode == HTP_OP_ARGSORT || node.opcode == HTP_OP_TOP_K) {
+            const auto * kparams = (const struct htp_sort_kernel_params *) node.kernel_params;
+            snprintf(str, max_size, "%s nth %d nchk %d chk %d vtcm %d",
+                     node.opcode == HTP_OP_TOP_K ? "top_k" : "argsort",
+                     (int) kparams->n_threads, (int) kparams->n_chunks,
+                     (int) kparams->chunk_elems, (int) kparams->vtcm_size);
+        } else if (node.opcode == HTP_OP_GET_ROWS) {
+            const auto * kparams = (const struct htp_get_rows_kernel_params *) node.kernel_params;
+            const char * ktype_str = "unknown";
+            switch (kparams->kernel_type) {
+                case HTP_GET_ROWS_KERNEL_SAMETYPE: ktype_str = "sametype"; break;
+                case HTP_GET_ROWS_KERNEL_TILED:    ktype_str = "tiled"; break;
+                case HTP_GET_ROWS_KERNEL_FLAT:     ktype_str = "flat"; break;
+            }
+            snprintf(str, max_size, "%s%s vtcm %u", ktype_str, kparams->n_threads > 1 ? "-multi" : "", (unsigned int) kparams->vtcm_size);
         } else {
             snprintf(str, max_size, "----");
         }
+
     }
 
     void format(const htp_opnode & node) {

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

@@ -0,0 +1,178 @@
+#ifndef HTP_ARGSORT_OPS_H
+#define HTP_ARGSORT_OPS_H
+
+#include <stdint.h>
+#include <stddef.h>
+#include <stdbool.h>
+#include <string.h>
+
+#include "hex-fastdiv.h"
+
+struct htp_sort_kernel_params {
+    int32_t  n_threads;
+    int32_t  total_rows;
+    int32_t  row_start;
+    int32_t  row_end;
+    int32_t  ne00;
+    int32_t  k;
+    int32_t  order;             // GGML_SORT_ORDER_ASC (0) or GGML_SORT_ORDER_DESC (1)
+    int32_t  is_top_k;          // 1 if TOP_K, 0 if ARGSORT
+    int32_t  use_dma;           // 1 if DMA enabled
+    int32_t  chunk_elems;
+    int32_t  n_chunks;
+    int32_t  vtcm_size;
+    int32_t  phase1_slot_size;
+    int32_t  merge_values_off;
+    int32_t  merge_indices_off;
+    int32_t  merge_elems;
+    int32_t  n_slots;
+    int32_t  pad[15];
+};
+
+struct htp_sort_vtcm_layout {
+    size_t   total_bytes;
+    size_t   phase1_slot_size;
+    size_t   merge_values_off;
+    size_t   merge_indices_off;
+    uint32_t chunk_elems;
+    uint32_t n_chunks;
+    uint32_t merge_elems;
+    uint32_t n_threads;
+    uint32_t n_slots;
+};
+
+static inline bool htp_sort_solve_layout(
+    struct htp_sort_vtcm_layout * layout,
+    uint32_t ne00,
+    uint32_t total_rows,
+    uint32_t k,
+    uint32_t n_threads_max,
+    size_t vtcm_budget,
+    bool is_top_k) {
+
+    memset(layout, 0, sizeof(*layout));
+
+    if (total_rows > 1) {
+        uint32_t n_threads = total_rows < n_threads_max ? total_rows : n_threads_max;
+
+        uint32_t n_vec = (ne00 + 31) / 32;
+        uint32_t n_vec_pow2 = 1;
+        while (n_vec_pow2 < n_vec) n_vec_pow2 <<= 1;
+        uint32_t ne00_padded = n_vec_pow2 * 32;
+
+        size_t values_size  = ((ne00_padded * sizeof(float)) + 127) & ~127;
+        size_t indices_size = ((ne00_padded * sizeof(int32_t)) + 127) & ~127;
+        size_t spad_per_slot = ((values_size + indices_size) + 255) & ~255;
+
+        uint32_t n_slots = 2;
+        if (spad_per_slot * 2 > vtcm_budget) {
+            n_slots = 1;
+        }
+
+        size_t spad_per_thread = spad_per_slot * n_slots;
+
+        while (n_threads > 1 && (spad_per_thread * n_threads) > vtcm_budget) {
+            n_threads--;
+        }
+
+        size_t total_bytes = spad_per_thread * n_threads;
+        if (total_bytes > vtcm_budget) {
+            return false;
+        }
+
+        layout->total_bytes      = total_bytes;
+        layout->phase1_slot_size = spad_per_slot;
+        layout->chunk_elems      = ne00_padded;
+        layout->n_chunks         = 1;
+        layout->n_threads        = n_threads;
+        layout->n_slots          = n_slots;
+        return true;
+    }
+
+    uint32_t n_vec = (ne00 + 31) / 32;
+    uint32_t n_vec_pow2 = 1;
+    while (n_vec_pow2 < n_vec) n_vec_pow2 <<= 1;
+
+    uint32_t n_chunks = 1;
+    if (ne00 > 1024) {
+        while (n_chunks * 2 <= n_threads_max && n_chunks * 2 <= n_vec_pow2) {
+            n_chunks *= 2;
+        }
+    }
+
+    uint32_t chunk_n_vec = n_vec_pow2 / n_chunks;
+    uint32_t chunk_elems = chunk_n_vec * 32;
+
+    size_t phase1_values_size  = ((chunk_elems * sizeof(float)) + 127) & ~127;
+    size_t phase1_indices_size = ((chunk_elems * sizeof(int32_t)) + 127) & ~127;
+    size_t phase1_slot_size    = ((phase1_values_size + phase1_indices_size) + 255) & ~255;
+    size_t phase1_total_size   = phase1_slot_size * n_chunks;
+
+    size_t merge_values_size  = 0;
+    size_t merge_indices_size = 0;
+    size_t merge_values_off   = phase1_total_size;
+    size_t merge_indices_off  = 0;
+    uint32_t merge_elems      = 0;
+
+    if (n_chunks > 1) {
+        if (is_top_k) {
+            uint32_t local_k = k < chunk_elems ? k : chunk_elems;
+            uint32_t total_candidates = n_chunks * local_k;
+            uint32_t merge_n_vec = (total_candidates + 31) / 32;
+            uint32_t merge_n_vec_pow2 = 1;
+            while (merge_n_vec_pow2 < merge_n_vec) merge_n_vec_pow2 <<= 1;
+            merge_elems = merge_n_vec_pow2 * 32;
+        } else {
+            merge_elems = n_vec_pow2 * 32;
+        }
+        merge_values_size  = ((merge_elems * sizeof(float)) + 127) & ~127;
+        merge_indices_size = ((merge_elems * sizeof(int32_t)) + 127) & ~127;
+        merge_indices_off  = merge_values_off + merge_values_size;
+    }
+
+    size_t total_bytes = phase1_total_size + merge_values_size + merge_indices_size;
+
+    if (total_bytes > vtcm_budget && n_chunks > 1) {
+        n_chunks = 1;
+        chunk_elems = n_vec_pow2 * 32;
+        phase1_values_size  = ((chunk_elems * sizeof(float)) + 127) & ~127;
+        phase1_indices_size = ((chunk_elems * sizeof(int32_t)) + 127) & ~127;
+        phase1_slot_size    = ((phase1_values_size + phase1_indices_size) + 255) & ~255;
+        phase1_total_size   = phase1_slot_size;
+        merge_values_size   = 0;
+        merge_indices_size  = 0;
+        merge_values_off    = phase1_total_size;
+        merge_indices_off   = 0;
+        merge_elems         = 0;
+        total_bytes         = phase1_total_size;
+    }
+
+    if (total_bytes > vtcm_budget) {
+        return false;
+    }
+
+    uint32_t n_slots = 1;
+    if (n_chunks == 1 && phase1_slot_size * 2 <= vtcm_budget) {
+        n_slots = 2;
+        total_bytes = phase1_slot_size * 2;
+    }
+
+    layout->total_bytes       = total_bytes;
+    layout->phase1_slot_size  = phase1_slot_size;
+    layout->merge_values_off  = merge_values_off;
+    layout->merge_indices_off = merge_indices_off;
+    layout->chunk_elems       = chunk_elems;
+    layout->n_chunks          = n_chunks;
+    layout->merge_elems       = merge_elems;
+    layout->n_threads         = n_chunks;
+    layout->n_slots           = n_slots;
+    return true;
+}
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_sort_kernel_params) <= 128, "htp_sort_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_sort_kernel_params) <= 128, "htp_sort_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_ARGSORT_OPS_H

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

@@ -917,6 +917,154 @@ static void binary_thread_add_id_f32(unsigned int nth, unsigned int ith, void *
     dma_queue_flush(dma_q);
 }
 
+static inline void hvx_div_scalar_f32(uint8_t * restrict dst, const uint8_t * restrict src, const float val, const uint32_t num_elems) {
+    hvx_mul_scalar_f32(dst, src, 1.0f / val, num_elems);
+}
+
+static inline void hvx_div_scalar_f16(uint8_t * restrict dst, const uint8_t * restrict src, const _Float16 val, const uint32_t num_elems) {
+    hvx_div_scalar_f16_aa(dst, src, val, num_elems);
+}
+
+typedef void (*compute_binary_chunked_t)(
+    uint8_t * restrict dst,
+    const uint8_t * restrict src0,
+    const uint8_t * restrict src1,
+    const uint32_t num_elems
+);
+
+typedef void (*compute_binary_scalar_chunked_f32_t)(
+    uint8_t * restrict dst,
+    const uint8_t * restrict src0,
+    const float val,
+    const uint32_t num_elems
+);
+
+typedef void (*compute_binary_scalar_chunked_f16_t)(
+    uint8_t * restrict dst,
+    const uint8_t * restrict src0,
+    const _Float16 val,
+    const uint32_t num_elems
+);
+
+struct binary_chunked_context {
+    struct htp_ops_context * octx;
+    struct htp_binary_vtcm_layout vtcm_layout;
+    uint8_t *                vtcm_base;
+    uint32_t                 elem_start;
+    uint32_t                 nelem;
+    uint32_t                 chunk_size;
+    uint32_t                 chunks_per_thread;
+    uint32_t                 total_chunks;
+    bool                     is_scalar;
+    float                    scalar_f32;
+    _Float16                 scalar_f16;
+    compute_binary_chunked_t compute;
+    compute_binary_scalar_chunked_f32_t compute_scalar_f32;
+    compute_binary_scalar_chunked_f16_t compute_scalar_f16;
+};
+
+static void binary_thread_chunked(unsigned int nth, unsigned int ith, void * data) {
+    (void) nth;
+    struct binary_chunked_context * ctx = (struct binary_chunked_context *) data;
+    struct htp_ops_context * octx = ctx->octx;
+    const struct htp_tensor * src0 = octx->src[0];
+    const struct htp_tensor * src1 = octx->src[1];
+    const struct htp_tensor * dst  = octx->dst;
+
+    const uint32_t start_chunk = ctx->chunks_per_thread * ith;
+    const uint32_t end_chunk   = MIN(start_chunk + ctx->chunks_per_thread, ctx->total_chunks);
+    if (start_chunk >= end_chunk) {
+        return;
+    }
+
+    const uint32_t src0_type   = src0->type;
+    const size_t   elem_size   = (src0_type == HTP_TYPE_F32) ? sizeof(float) : sizeof(_Float16);
+    const uint32_t chunk_size  = ctx->chunk_size;
+    const size_t   chunk_bytes = ctx->vtcm_layout.src0_spad_half_size;
+
+    FARF(HIGH, "binary-chunked: %d/%d (%u:%u) chunks %u elems %u",
+         ith, nth, start_chunk, end_chunk, ctx->total_chunks, ctx->nelem);
+
+    const struct htp_binary_vtcm_layout * layout = &ctx->vtcm_layout;
+    uint8_t * src0_spad_base = VTCM_LAYOUT_PTR(uint8_t, ctx->vtcm_base, layout->off_src0) + (ith * layout->src0_bytes_per_thread);
+    uint8_t * src1_spad_base = ctx->is_scalar ? NULL : (VTCM_LAYOUT_PTR(uint8_t, ctx->vtcm_base, layout->off_src1) + (ith * layout->src1_bytes_per_thread));
+    uint8_t * dst_spad_base  = VTCM_LAYOUT_PTR(uint8_t, ctx->vtcm_base, layout->off_dst)  + (ith * layout->dst_bytes_per_thread);
+
+    dma_queue * dma_q = octx->ctx->dma[ith];
+    uint32_t prefetch_chunk = start_chunk;
+    int spad_idx = 0;
+
+    for (int k = 0; k < 2 && prefetch_chunk < end_chunk; k++) {
+        const uint32_t c_start   = ctx->elem_start + prefetch_chunk * chunk_size;
+        const uint32_t c_end     = MIN(c_start + chunk_size, ctx->elem_start + ctx->nelem);
+        const uint32_t cur_elems = c_end - c_start;
+        const uint32_t cur_bytes = cur_elems * elem_size;
+
+        dma_addr_t s0_curr = src0->data + (size_t) c_start * elem_size;
+        dma_addr_t d_curr  = dst->data  + (size_t) c_start * elem_size;
+
+        uint8_t * s0_spad = src0_spad_base + spad_idx * chunk_bytes;
+        uint8_t * d_spad  = dst_spad_base  + spad_idx * chunk_bytes;
+
+        dma_queue_push(dma_q, dma_make_data(d_curr, d_spad), chunk_bytes, chunk_bytes, cur_bytes, 0);
+        dma_queue_push(dma_q, dma_make_data(s0_spad, s0_curr), chunk_bytes, chunk_bytes, cur_bytes, 1);
+        if (!ctx->is_scalar) {
+            dma_addr_t s1_curr = src1->data + (size_t) c_start * elem_size;
+            uint8_t * s1_spad = src1_spad_base + spad_idx * chunk_bytes;
+            dma_queue_push(dma_q, dma_make_data(s1_spad, s1_curr), chunk_bytes, chunk_bytes, cur_bytes, 1);
+        }
+
+        prefetch_chunk++;
+        spad_idx ^= 1;
+    }
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
+    for (uint32_t c = start_chunk; c < end_chunk; c++) {
+        const uint32_t c_start   = ctx->elem_start + c * chunk_size;
+        const uint32_t c_end     = MIN(c_start + chunk_size, ctx->elem_start + ctx->nelem);
+        const uint32_t cur_elems = c_end - c_start;
+        const uint32_t cur_bytes = cur_elems * elem_size;
+
+        uint8_t * d_spad  = (uint8_t *) dma_queue_pop(dma_q).src;
+        uint8_t * s0_spad = (uint8_t *) dma_queue_pop(dma_q).dst;
+        uint8_t * s1_spad = ctx->is_scalar ? NULL : (uint8_t *) dma_queue_pop(dma_q).dst;
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) c);
+        if (ctx->is_scalar) {
+            if (src0_type == HTP_TYPE_F32) {
+                ctx->compute_scalar_f32(d_spad, s0_spad, ctx->scalar_f32, cur_elems);
+            } else {
+                ctx->compute_scalar_f16(d_spad, s0_spad, ctx->scalar_f16, cur_elems);
+            }
+        } else {
+            ctx->compute(d_spad, s0_spad, s1_spad, cur_elems);
+        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) c);
+
+        dma_addr_t dst_curr = dst->data + (size_t) c_start * elem_size;
+        dma_queue_push(dma_q, dma_make_data(dst_curr, d_spad), chunk_bytes, chunk_bytes, cur_bytes, 1);
+
+        if (prefetch_chunk < end_chunk) {
+            const uint32_t pc_start  = ctx->elem_start + prefetch_chunk * chunk_size;
+            const uint32_t pc_end    = MIN(pc_start + chunk_size, ctx->elem_start + ctx->nelem);
+            const uint32_t p_elems   = pc_end - pc_start;
+            const uint32_t p_bytes   = p_elems * elem_size;
+
+            dma_addr_t s0_next = src0->data + (size_t) pc_start * elem_size;
+            dma_queue_push(dma_q, dma_make_data(s0_spad, s0_next), chunk_bytes, chunk_bytes, p_bytes, 1);
+            if (!ctx->is_scalar) {
+                dma_addr_t s1_next = src1->data + (size_t) pc_start * elem_size;
+                dma_queue_push(dma_q, dma_make_data(s1_spad, s1_next), chunk_bytes, chunk_bytes, p_bytes, 1);
+            }
+
+            prefetch_chunk++;
+        }
+    }
+
+    dma_queue_flush(dma_q);
+}
+
 static int execute_op_binary(struct htp_ops_context * octx) {
     const struct htp_tensor * src0 = octx->src[0];
     const struct htp_tensor * src1 = octx->src[1];
@@ -932,6 +1080,135 @@ static int execute_op_binary(struct htp_ops_context * octx) {
     const size_t src1_row_size = src1->ne[0] * elem_size;
     const size_t dst_row_size  = dst->ne[0]  * elem_size;
 
+    if (kparams->kernel_type == HTP_BINARY_KERNEL_CHUNKED) {
+        const bool is_scalar = kparams->is_scalar || (src1->ne[0] == 1 && src1->ne[1] == 1 && src1->ne[2] == 1 && src1->ne[3] == 1);
+
+        compute_binary_chunked_t compute = NULL;
+        compute_binary_scalar_chunked_f32_t compute_scalar_f32 = NULL;
+        compute_binary_scalar_chunked_f16_t compute_scalar_f16 = NULL;
+
+        if (is_scalar) {
+            if (src0_type == HTP_TYPE_F32) {
+                switch (octx->op) {
+                    case HTP_OP_ADD: compute_scalar_f32 = hvx_add_scalar_f32; break;
+                    case HTP_OP_SUB: compute_scalar_f32 = hvx_sub_scalar_f32; break;
+                    case HTP_OP_MUL: compute_scalar_f32 = hvx_mul_scalar_f32; break;
+                    case HTP_OP_DIV: compute_scalar_f32 = hvx_div_scalar_f32; break;
+                    default: break;
+                }
+            } else if (src0_type == HTP_TYPE_F16) {
+                switch (octx->op) {
+                    case HTP_OP_ADD: compute_scalar_f16 = hvx_add_scalar_f16; break;
+                    case HTP_OP_SUB: compute_scalar_f16 = hvx_sub_scalar_f16; break;
+                    case HTP_OP_MUL: compute_scalar_f16 = hvx_mul_scalar_f16; break;
+                    case HTP_OP_DIV: compute_scalar_f16 = hvx_div_scalar_f16; break;
+                    default: break;
+                }
+            }
+            if (!compute_scalar_f32 && !compute_scalar_f16) {
+                return HTP_STATUS_NO_SUPPORT;
+            }
+        } else {
+            if (src0_type == HTP_TYPE_F32) {
+                switch (octx->op) {
+                    case HTP_OP_ADD: compute = hvx_add_f32; break;
+                    case HTP_OP_SUB: compute = hvx_sub_f32; break;
+                    case HTP_OP_MUL: compute = hvx_mul_f32; break;
+                    case HTP_OP_DIV: compute = hvx_div_f32; break;
+                    default: break;
+                }
+            } else if (src0_type == HTP_TYPE_F16) {
+                switch (octx->op) {
+                    case HTP_OP_ADD: compute = hvx_add_f16; break;
+                    case HTP_OP_SUB: compute = hvx_sub_f16; break;
+                    case HTP_OP_MUL: compute = hvx_mul_f16; break;
+                    case HTP_OP_DIV: compute = hvx_div_f16; break;
+                    default: break;
+                }
+            }
+            if (!compute) {
+                return HTP_STATUS_NO_SUPPORT;
+            }
+        }
+
+        const uint32_t total_elems = (uint32_t) (src0->ne[0] * src0->ne[1] * src0->ne[2] * src0->ne[3]);
+        if (total_elems == 0) {
+            return HTP_STATUS_OK;
+        }
+
+        uint32_t elem_start = 0;
+        uint32_t nelem      = total_elems;
+        const uint32_t elems_per_line = (src0_type == HTP_TYPE_F32) ? 32 : 64;
+
+        if (octx->ctx->mdev.count > 1) {
+            const bool can_split = htp_tensor_mdev_data_aligned(dst) &&
+                                   htp_tensor_mdev_data_aligned(src0) &&
+                                   (is_scalar || htp_tensor_mdev_data_aligned(src1)) &&
+                                   htp_tensor_is_contiguous(dst, elem_size) &&
+                                   htp_tensor_is_contiguous(src0, elem_size) &&
+                                   (is_scalar || htp_tensor_is_contiguous(src1, elem_size));
+            const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(
+                total_elems, can_split ? elems_per_line : 0,
+                octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+            elem_start = range.start;
+            nelem      = range.count;
+        }
+
+        if (nelem == 0) {
+            return HTP_STATUS_OK;
+        }
+
+        if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+            return HTP_STATUS_INVAL_PARAMS;
+        }
+
+        struct htp_binary_vtcm_layout vtcm_layout;
+        htp_binary_vtcm_layout_build(&vtcm_layout, kparams, octx->ctx->vtcm_size);
+        if (vtcm_layout.total_bytes == 0 || vtcm_layout.total_bytes > octx->ctx->vtcm_size) {
+            return HTP_STATUS_VTCM_TOO_SMALL;
+        }
+
+        const uint32_t chunk_size = kparams->chunk_size > 0 ? kparams->chunk_size : (32768 / elem_size);
+        const uint32_t total_chunks = (nelem + chunk_size - 1) / chunk_size;
+        const uint32_t n_threads = (total_chunks >= 2) ? MIN(octx->n_threads, total_chunks) : 1;
+        const uint32_t chunks_per_thread = (total_chunks + n_threads - 1) / n_threads;
+
+        float scalar_f32 = 0.0f;
+        _Float16 scalar_f16 = 0;
+        if (is_scalar) {
+            uint8_t * vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, octx->ctx->vtcm_base, vtcm_layout.off_src1);
+            dma_queue * dma_q = octx->ctx->dma[0];
+            dma_queue_push(dma_q, dma_make_data(vtcm_src1, src1->data), 128, 0, elem_size, 1);
+            dma_queue_pop(dma_q);
+
+            if (src0_type == HTP_TYPE_F32) {
+                scalar_f32 = ((const float *) vtcm_src1)[0];
+            } else {
+                scalar_f16 = ((const _Float16 *) vtcm_src1)[0];
+            }
+        }
+
+        struct binary_chunked_context cctx = {
+            .octx                = octx,
+            .vtcm_layout         = vtcm_layout,
+            .vtcm_base           = (uint8_t *) octx->ctx->vtcm_base,
+            .elem_start          = elem_start,
+            .nelem               = nelem,
+            .chunk_size          = chunk_size,
+            .chunks_per_thread   = chunks_per_thread,
+            .total_chunks        = total_chunks,
+            .is_scalar           = is_scalar,
+            .scalar_f32          = scalar_f32,
+            .scalar_f16          = scalar_f16,
+            .compute             = compute,
+            .compute_scalar_f32  = compute_scalar_f32,
+            .compute_scalar_f16  = compute_scalar_f16,
+        };
+
+        work_queue_run(octx->ctx->work_queue, binary_thread_chunked, &cctx, n_threads);
+        return HTP_STATUS_OK;
+    }
+
     uint32_t row_start = 0;
     uint32_t nrows     = src0_nrows;
 

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

@@ -17,6 +17,7 @@ enum htp_binary_kernel_type {
     HTP_BINARY_KERNEL_ADD_ID,
     HTP_BINARY_KERNEL_COMPLEX,
     HTP_BINARY_KERNEL_REPEAT,
+    HTP_BINARY_KERNEL_CHUNKED,
 };
 
 struct htp_binary_kernel_params {
@@ -30,6 +31,10 @@ struct htp_binary_kernel_params {
 
     uint32_t src1_size;
     uint32_t vtcm_size;
+
+    uint32_t chunk_size;
+    uint32_t chunk_bytes;
+    uint32_t is_scalar;
 };
 
 #if defined(__cplusplus)
@@ -68,6 +73,40 @@ static inline void htp_binary_vtcm_layout_build(
         return;
     }
 
+    if (kparams->kernel_type == HTP_BINARY_KERNEL_CHUNKED) {
+        const size_t chunk_bytes = kparams->chunk_bytes;
+        if (chunk_bytes == 0) {
+            return;
+        }
+
+        L->src0_bytes_per_thread = 2 * chunk_bytes;
+        L->src1_bytes_per_thread = kparams->is_scalar ? 0 : (2 * chunk_bytes);
+        L->dst_bytes_per_thread  = 2 * chunk_bytes;
+
+        L->src0_spad_half_size = chunk_bytes;
+        L->src1_spad_half_size = kparams->is_scalar ? 0 : chunk_bytes;
+        L->dst_spad_half_size  = chunk_bytes;
+
+        L->rows_per_buffer = 1;
+        L->src1_size = 0;
+
+        const size_t src0_total = n_threads * L->src0_bytes_per_thread;
+        const size_t src1_total = kparams->is_scalar ? 128 : (n_threads * L->src1_bytes_per_thread);
+        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_dst,  dst_total);
+
+        if (off > vtcm_size) {
+            return;
+        }
+
+        L->total_bytes = off;
+        return;
+    }
+
     const size_t spad_row_total = (kparams->kernel_type == HTP_BINARY_KERNEL_SAME_SHAPE)
         ? 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);

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

@@ -301,6 +301,86 @@ static void cpy_thread_f32_f16_sameshape(unsigned int nth, unsigned int ith, voi
     }
 }
 
+static void cpy_thread_i32_f32_sameshape(unsigned int nth, unsigned int ith, void * data) {
+    struct htp_copy_context * ct = (struct htp_copy_context *) data;
+    struct htp_ops_context * octx = ct->octx;
+    cpy_preamble;
+
+    const uint32_t dr  = ct->src0_nrows_per_thread;
+    const uint32_t ir0 = ct->row_start + dr * ith;
+    const uint32_t ir1 = MIN(ir0 + dr, ct->row_start + ct->nrows);
+    if (ir0 >= ir1) return;
+
+    const uint32_t ne02_ne01 = ne02 * ne01;
+    uint32_t i03 = fastdiv(ir0, &ct->div_ne02_ne01);
+    uint32_t rem = ir0 - i03 * ne02_ne01;
+    uint32_t i02 = fastdiv(rem, &ct->div_ne01);
+    uint32_t i01 = rem - i02 * ne01;
+
+    uint8_t* dst_ptr  = (uint8_t*) dst->data  + i01*nb1  + i02*nb2  + i03*nb3;
+    uint8_t* src0_ptr = (uint8_t*) src0->data + i01*nb01 + i02*nb02 + i03*nb03;
+
+    for (uint32_t r = ir0; r < ir1; r++) {
+        hex_l2fetch(src0_ptr, ne00 * sizeof(float), nb01, 2);
+        const float * restrict src_row = (const float *) src0_ptr;
+        int32_t * restrict dst_row = (int32_t *) dst_ptr;
+        for (uint32_t i = 0; i < ne00; i++) {
+            dst_row[i] = (int32_t) src_row[i];
+        }
+        dst_ptr  += nb1;
+        src0_ptr += nb01;
+        if (++i01 == ne01) {
+            i01 = 0;
+            if (++i02 == ne02) {
+                i02 = 0;
+                i03++;
+            }
+            dst_ptr  = (uint8_t*) dst->data  + i02*nb2  + i03*nb3;
+            src0_ptr = (uint8_t*) src0->data + i02*nb02 + i03*nb03;
+        }
+    }
+}
+
+static void cpy_thread_f32_i32_sameshape(unsigned int nth, unsigned int ith, void * data) {
+    struct htp_copy_context * ct = (struct htp_copy_context *) data;
+    struct htp_ops_context * octx = ct->octx;
+    cpy_preamble;
+
+    const uint32_t dr  = ct->src0_nrows_per_thread;
+    const uint32_t ir0 = ct->row_start + dr * ith;
+    const uint32_t ir1 = MIN(ir0 + dr, ct->row_start + ct->nrows);
+    if (ir0 >= ir1) return;
+
+    const uint32_t ne02_ne01 = ne02 * ne01;
+    uint32_t i03 = fastdiv(ir0, &ct->div_ne02_ne01);
+    uint32_t rem = ir0 - i03 * ne02_ne01;
+    uint32_t i02 = fastdiv(rem, &ct->div_ne01);
+    uint32_t i01 = rem - i02 * ne01;
+
+    uint8_t* dst_ptr  = (uint8_t*) dst->data  + i01*nb1  + i02*nb2  + i03*nb3;
+    uint8_t* src0_ptr = (uint8_t*) src0->data + i01*nb01 + i02*nb02 + i03*nb03;
+
+    for (uint32_t r = ir0; r < ir1; r++) {
+        hex_l2fetch(src0_ptr, ne00 * sizeof(int32_t), nb01, 2);
+        const int32_t * restrict src_row = (const int32_t *) src0_ptr;
+        float * restrict dst_row = (float *) dst_ptr;
+        for (uint32_t i = 0; i < ne00; i++) {
+            dst_row[i] = (float) src_row[i];
+        }
+        dst_ptr  += nb1;
+        src0_ptr += nb01;
+        if (++i01 == ne01) {
+            i01 = 0;
+            if (++i02 == ne02) {
+                i02 = 0;
+                i03++;
+            }
+            dst_ptr  = (uint8_t*) dst->data  + i02*nb2  + i03*nb3;
+            src0_ptr = (uint8_t*) src0->data + i02*nb02 + i03*nb03;
+        }
+    }
+}
+
 static inline void cpy_dma_push_2d_chunked(
     dma_queue * dma_q,
     dma_addr_t  dst,
@@ -366,6 +446,42 @@ static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) {
     cpy_preamble;
     *use_dma = false;
 
+    const uint32_t total_elems_src = ne00 * ne01 * ne02 * ne03;
+    const uint32_t total_elems_dst = ne0 * ne1 * ne2 * ne3;
+    if (total_elems_src == 1 && total_elems_dst == 1) {
+        if (octx->ctx->mdev.count > 1 && octx->ctx->mdev.idx > 0) {
+            return HTP_STATUS_OK;
+        }
+        if (src0->type == HTP_TYPE_F32 && dst->type == HTP_TYPE_I32) {
+            ((int32_t *) dst->data)[0] = (int32_t) (((const float *) src0->data)[0]);
+            return HTP_STATUS_OK;
+        }
+        if (src0->type == HTP_TYPE_I32 && dst->type == HTP_TYPE_F32) {
+            ((float *) dst->data)[0] = (float) (((const int32_t *) src0->data)[0]);
+            return HTP_STATUS_OK;
+        }
+        if (src0->type == HTP_TYPE_I32 && dst->type == HTP_TYPE_I32) {
+            ((int32_t *) dst->data)[0] = ((const int32_t *) src0->data)[0];
+            return HTP_STATUS_OK;
+        }
+        if (src0->type == HTP_TYPE_F32 && dst->type == HTP_TYPE_F32) {
+            ((float *) dst->data)[0] = ((const float *) src0->data)[0];
+            return HTP_STATUS_OK;
+        }
+        if (src0->type == HTP_TYPE_F16 && dst->type == HTP_TYPE_F16) {
+            ((__fp16 *) dst->data)[0] = ((const __fp16 *) src0->data)[0];
+            return HTP_STATUS_OK;
+        }
+        if (src0->type == HTP_TYPE_F32 && dst->type == HTP_TYPE_F16) {
+            ((__fp16 *) dst->data)[0] = (__fp16) (((const float *) src0->data)[0]);
+            return HTP_STATUS_OK;
+        }
+        if (src0->type == HTP_TYPE_F16 && dst->type == HTP_TYPE_F32) {
+            ((float *) dst->data)[0] = (float) (((const __fp16 *) src0->data)[0]);
+            return HTP_STATUS_OK;
+        }
+    }
+
     struct htp_copy_context ct;
     ct.octx = octx;
 
@@ -450,6 +566,10 @@ static int exec_cpy(struct htp_ops_context * octx, bool * use_dma) {
                 copy_fun = cpy_thread_f16_f32_sameshape;
             } else if (dst->type == HTP_TYPE_F32 && src0->type == HTP_TYPE_F16) {
                 copy_fun = cpy_thread_f32_f16_sameshape;
+            } else if (dst->type == HTP_TYPE_I32 && src0->type == HTP_TYPE_F32) {
+                copy_fun = cpy_thread_i32_f32_sameshape;
+            } else if (dst->type == HTP_TYPE_F32 && src0->type == HTP_TYPE_I32) {
+                copy_fun = cpy_thread_f32_i32_sameshape;
             } else {
                 return HTP_STATUS_NO_SUPPORT;
             }

ggml/src/ggml-hexagon/htp/dma-queue.h

@@ -426,6 +426,22 @@ static inline bool dma_queue_push(dma_queue *q, dma_data ddata, size_t dst_strid
 
 #endif
 
+static inline void dma_sync_read(dma_queue * dma_q, void * dst, dma_addr_t src, size_t bytes) {
+    const uint32_t b = (uint32_t) bytes;
+    if (b > 0) {
+        dma_queue_push(dma_q, dma_make_data(dst, src), b, b, b, 1);
+        dma_queue_pop(dma_q);
+    }
+}
+
+static inline void dma_sync_write(dma_queue * dma_q, dma_addr_t dst, const void * src, size_t bytes) {
+    const uint32_t b = (uint32_t) bytes;
+    if (b > 0) {
+        dma_queue_push(dma_q, dma_make_data(dst, src), b, b, b, 1);
+        dma_queue_pop(dma_q);
+    }
+}
+
 #define DMA_CACHE_MAX_SIZE 256U
 
 // Fully assoc LRU cache

ggml/src/ggml-hexagon/htp/get-rows-ops.c

@@ -17,6 +17,7 @@
 #include "htp-tensor.h"
 #include "hvx-utils.h"
 #include "hvx-quant.h"
+#include "matmul-ops.h"
 #include "get-rows-ops.h"
 #include "work-queue.h"
 
@@ -28,6 +29,9 @@ struct get_rows_context {
     uint32_t task_start;
     uint32_t tasks;
     uint32_t tasks_per_thread;
+    uint32_t tile_size;
+    uint32_t tile_stride;
+    bool index_i32;
 };
 
 #define get_rows_preamble                      \
@@ -195,32 +199,191 @@ static void get_rows_thread_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsigned
     dma_queue_flush(dma_q);                                                                                                     \
 }
 
-#define F32_BYTES(n)  ((n) * sizeof(float))
 #define F16_BYTES(n)  ((n) * sizeof(__fp16))
 #define Q8_0_BYTES(n) (((n) / 32) * sizeof(block_q8_0))
 
-GET_ROWS_THREAD_DT_FN(f32,  F32_BYTES,  int32_t, { if (cur_elems > 0) hvx_copy_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, cur_elems); })
-GET_ROWS_THREAD_DT_FN(f32,  F32_BYTES,  int64_t, { if (cur_elems > 0) hvx_copy_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, cur_elems); })
+static __attribute__((noinline)) void compute_get_rows_f16(float * dst_spad, const void * src_spad, uint32_t cur_elems) {
+    hvx_dequantize_row_f16_f32(dst_spad, src_spad, cur_elems);
+}
 
-GET_ROWS_THREAD_DT_FN(f16,  F16_BYTES,  int32_t, { hvx_dequantize_row_f16_f32((float *)dst_spad, src_spad, ne00); })
-GET_ROWS_THREAD_DT_FN(f16,  F16_BYTES,  int64_t, { hvx_dequantize_row_f16_f32((float *)dst_spad, src_spad, ne00); })
+static __attribute__((noinline)) void compute_get_rows_q8_0(float * dst_spad, const void * src_spad, uint32_t cur_elems) {
+    hvx_dequantize_row_q8_0_f32(dst_spad, src_spad, cur_elems);
+}
 
-GET_ROWS_THREAD_DT_FN(q8_0, Q8_0_BYTES, int32_t, { hvx_dequantize_row_q8_0_f32((float *)dst_spad, src_spad, ne00); })
-GET_ROWS_THREAD_DT_FN(q8_0, Q8_0_BYTES, int64_t, { hvx_dequantize_row_q8_0_f32((float *)dst_spad, src_spad, ne00); })
+GET_ROWS_THREAD_DT_FN(f16,  F16_BYTES,  int32_t, { compute_get_rows_f16((float *)dst_spad, src_spad, cur_elems); })
+GET_ROWS_THREAD_DT_FN(f16,  F16_BYTES,  int64_t, { compute_get_rows_f16((float *)dst_spad, src_spad, cur_elems); })
+
+GET_ROWS_THREAD_DT_FN(q8_0, Q8_0_BYTES, int32_t, { compute_get_rows_q8_0((float *)dst_spad, src_spad, cur_elems); })
+GET_ROWS_THREAD_DT_FN(q8_0, Q8_0_BYTES, int64_t, { compute_get_rows_q8_0((float *)dst_spad, src_spad, cur_elems); })
+
+
+static __attribute__((noinline)) void compute_get_rows_tiled(float * dst, const uint8_t * tile, uint32_t row, bool q4) {
+    const HVX_VectorPred first2 = Q6_Q_vsetq_R(2);
+    const HVX_VectorPred first4 = Q6_Q_vsetq_R(4);
+    HVX_Vector vq = Q6_V_vzero();
+    if (q4) {
+        const HVX_VectorPred first1 = Q6_Q_vsetq_R(1);
+        const HVX_VectorPred first3 = Q6_Q_vsetq_R(3);
+        for (int group = 3; group >= 0; --group) {
+            const HVX_Vector v = Q6_V_vror_VR(hvx_vmem(tile + group * VLEN), row);
+            // Four planes contribute bytes at 0, 32, 64 and 96 after rotation.
+            HVX_Vector packed = Q6_V_vmux_QVV(first1, v, Q6_V_vror_VR(v, 31));
+            packed = Q6_V_vmux_QVV(first2, packed, Q6_V_vror_VR(v, 62));
+            packed = Q6_V_vmux_QVV(first3, packed, Q6_V_vror_VR(v, 93));
+            vq = Q6_V_vmux_QVV(first4, packed, Q6_V_vror_VR(vq, VLEN - 4));
+        }
+        const HVX_Vector lo = Q6_V_vand_VV(vq, Q6_Vb_vsplat_R(0x0F));
+        const HVX_Vector hi = Q6_Vub_vlsr_VubR(vq, 4);
+        vq = Q6_V_lo_W(Q6_W_vshuff_VVR(hi, lo, -1));
+        vq = Q6_Vb_vsub_VbVb(vq, Q6_Vb_vsplat_R(8));
+    } else {
+        for (int group = 7; group >= 0; --group) {
+            const HVX_Vector v = Q6_V_vror_VR(hvx_vmem(tile + group * VLEN), 2 * row);
+            // Two planes contribute halfwords at 0 and 64 after rotation.
+            const HVX_Vector packed = Q6_V_vmux_QVV(first2, v, Q6_V_vror_VR(v, 62));
+            vq = Q6_V_vmux_QVV(first4, packed, Q6_V_vror_VR(vq, VLEN - 4));
+        }
+    }
+    const HVX_Vector scales = hvx_vmem(tile + (q4 ? 512 : 1024));
+    const HVX_Vector scale_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 2 * row));
+    const HVX_Vector scale = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_hf));
+    const HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(vq);
+    const HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16));
+    const HVX_Vector values = hvx_vec_mul_f32_f32(Q6_Vsf_equals_Vw(Q6_V_lo_W(p32)), scale);
+    *(HVX_Vector *) dst = values;
+}
+
+struct get_rows_tiled_task {
+    dma_addr_t tile_src_base;
+    dma_addr_t dst_data;
+    uint32_t   row;
+};
+
+static inline struct get_rows_tiled_task get_rows_tiled_calc_task(
+    const struct htp_ops_context * octx,
+    const struct get_rows_context * grctx,
+    uint32_t i,
+    uint32_t n_k_tiles,
+    uint32_t tile_size
+) {
+    const struct htp_get_rows_kernel_params * kparams = grctx->kparams;
+    get_rows_preamble;
+
+    const uint32_t i12 = fastdiv(i, &kparams->div_ne10_ne11);
+    const uint32_t rem = i - i12 * ne11 * ne10;
+    const uint32_t i11 = fastdiv(rem, &kparams->div_ne10);
+    const uint32_t i10 = rem - i11 * ne10;
+    const dma_addr_t src1_data = octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12;
+    const uint32_t i01 = grctx->index_i32 ? *(const int32_t *)(uintptr_t) src1_data : (uint32_t) *(const int64_t *)(uintptr_t) src1_data;
+    assert(i01 < ne01);
+
+    const uint32_t q02 = fastdiv(i11, &kparams->div_ne02);
+    const uint32_t i02 = i11 - q02 * ne02;
+    const uint32_t q03 = fastdiv(i12, &kparams->div_ne03);
+    const uint32_t i03 = i12 - q03 * ne03;
+    const uint32_t column_tile = i01 / HTP_MM_HMX_TILE_N_ROWS;
+    const uint32_t row = i01 % HTP_MM_HMX_TILE_N_ROWS;
+    const dma_addr_t matrix = octx->src[0]->data + i02*nb02 + i03*nb03;
+
+    struct get_rows_tiled_task task;
+    task.tile_src_base = matrix + (column_tile * n_k_tiles) * tile_size;
+    task.dst_data      = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3;
+    task.row           = row;
+    return task;
+}
+
+static void get_rows_thread_tiled(unsigned int nth, unsigned int ith, void * data) {
+    struct get_rows_context * grctx = (struct get_rows_context *) data;
+    struct htp_ops_context * octx = grctx->octx;
+    const struct htp_get_rows_kernel_params * kparams = grctx->kparams;
+    get_rows_preamble;
+
+    const uint32_t dr  = grctx->tasks_per_thread;
+    const uint32_t ir0 = grctx->task_start + dr * ith;
+    if (ir0 >= grctx->task_start + grctx->tasks) {
+        return;
+    }
+
+    const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks);
+    const uint32_t n_k_tiles = ne00 / HTP_MM_HMX_TILE_N_COLS;
+    const struct htp_get_rows_vtcm_layout * vtcm_layout = &grctx->vtcm_layout;
+    uint8_t * src_spad_base = grctx->vtcm_base + vtcm_layout->off_src0 + ith * vtcm_layout->src0_bytes_per_thread;
+    uint8_t * dst_spad_base = grctx->vtcm_base + vtcm_layout->off_dst + ith * vtcm_layout->dst_bytes_per_thread;
+    dma_queue * dma_q = octx->ctx->dma[ith];
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
+    const uint32_t tile_size   = grctx->tile_size;
+    const uint32_t tile_stride = grctx->tile_stride;
+    const uint32_t dst_bytes   = ne00 * sizeof(float);
+    const bool is_q4 = (octx->src[0]->type == HTP_TYPE_Q4_0);
+
+    for (uint32_t step = 0, spad_idx = 0; step < ir1 - ir0 && spad_idx < 2; ++step, ++spad_idx) {
+        const uint32_t i = ir0 + step;
+        struct get_rows_tiled_task task = get_rows_tiled_calc_task(octx, grctx, i, n_k_tiles, tile_size);
+
+        // Dummy writeback to prime the queue with dst descriptor
+        dma_queue_push(dma_q,
+                       dma_make_data(task.dst_data, dst_spad_base + spad_idx * vtcm_layout->dst_spad_half_size),
+                       dst_bytes, vtcm_layout->dst_spad_half_size, dst_bytes, 0);
+
+        // Prefetch row tiles
+        dma_queue_push(dma_q,
+                       dma_make_data(src_spad_base + spad_idx * vtcm_layout->src0_spad_half_size, task.tile_src_base),
+                       tile_stride, tile_size, tile_size, n_k_tiles);
+    }
+
+    for (uint32_t step = 0; step < ir1 - ir0; ++step) {
+        const uint32_t i = ir0 + step;
+        float * dst_spad   = (float *) dma_queue_pop(dma_q).src;
+        uint8_t * src_spad = (uint8_t *) dma_queue_pop(dma_q).dst;
+
+        struct get_rows_tiled_task task = get_rows_tiled_calc_task(octx, grctx, i, n_k_tiles, tile_size);
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
+        for (uint32_t k_tile = 0; k_tile < n_k_tiles; ++k_tile) {
+            const uint8_t * tile = src_spad + k_tile * tile_stride;
+            float * dst_block = dst_spad + k_tile * HTP_MM_HMX_TILE_N_COLS;
+            compute_get_rows_tiled(dst_block, tile, task.row, is_q4);
+        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
+
+        // Real writeback of dst_spad
+        dma_queue_push(dma_q,
+                       dma_make_data(task.dst_data, dst_spad),
+                       dst_bytes, vtcm_layout->dst_spad_half_size, dst_bytes, 1);
+
+        const uint32_t next_step = step + 2;
+        if (next_step < ir1 - ir0) {
+            const uint32_t ni = ir0 + next_step;
+            struct get_rows_tiled_task next_task = get_rows_tiled_calc_task(octx, grctx, ni, n_k_tiles, tile_size);
+            dma_queue_push(dma_q,
+                           dma_make_data(src_spad, next_task.tile_src_base),
+                           tile_stride, tile_size, tile_size, n_k_tiles);
+        }
+    }
+
+    dma_queue_flush(dma_q);
+}
 
 int op_get_rows(struct htp_ops_context * octx) {
     const struct htp_get_rows_kernel_params * kparams = (const struct htp_get_rows_kernel_params *) octx->kernel_params;
 
     if (octx->src[0]->type != HTP_TYPE_F32 &&
-        octx->src[0]->type != HTP_TYPE_F16 &&
-        octx->src[0]->type != HTP_TYPE_Q8_0 &&
-        octx->src[0]->type != HTP_TYPE_I32) {
+         octx->src[0]->type != HTP_TYPE_F16 &&
+         octx->src[0]->type != HTP_TYPE_Q4_0 &&
+         octx->src[0]->type != HTP_TYPE_Q8_0 &&
+         octx->src[0]->type != HTP_TYPE_I32) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
-    if ((octx->src[0]->type == HTP_TYPE_I32 && octx->dst->type != HTP_TYPE_I32) ||
-        (octx->src[0]->type != HTP_TYPE_I32 && octx->dst->type != HTP_TYPE_F32)) {
-        return HTP_STATUS_NO_SUPPORT;
+    if (kparams->kernel_type == HTP_GET_ROWS_KERNEL_SAMETYPE) {
+        if (octx->src[0]->type != octx->dst->type) {
+            return HTP_STATUS_NO_SUPPORT;
+        }
+    } else {
+        if (octx->dst->type != HTP_TYPE_F32) {
+            return HTP_STATUS_NO_SUPPORT;
+        }
     }
 
     if (octx->src[1]->type != HTP_TYPE_I32 && octx->src[1]->type != HTP_TYPE_I64) {
@@ -262,33 +425,48 @@ int op_get_rows(struct htp_ops_context * octx) {
     grctx.vtcm_base = (uint8_t *)octx->ctx->vtcm_base;
     grctx.task_start = task_start;
     grctx.tasks = tasks;
-    grctx.tasks_per_thread = fastdiv(tasks + n_threads - 1, &octx->n_threads_div);
+    grctx.tasks_per_thread = octx->ctx->mdev.count == 1 ? kparams->tasks_per_thread : fastdiv(tasks + n_threads - 1, &octx->n_threads_div);
+    grctx.tile_size = octx->src[0]->type == HTP_TYPE_Q4_0 ? HTP_MM_WEIGHT_TILE_SIZE_Q4_0 : HTP_MM_WEIGHT_TILE_SIZE_Q8_0;
+    grctx.tile_stride = (grctx.tile_size + 127) & ~127;
+    grctx.index_i32 = octx->src[1]->type == HTP_TYPE_I32;
 
     const uint32_t ne00 = octx->src[0]->ne[0];
-    htp_get_rows_vtcm_layout_build(&grctx.vtcm_layout, octx->src[0]->type, ne00, n_threads);
+    htp_get_rows_vtcm_layout_build(&grctx.vtcm_layout, kparams->kernel_type, octx->src[0]->type, ne00, n_threads);
+
+    if (grctx.vtcm_layout.total_bytes > octx->ctx->vtcm_size) {
+        FARF(ERROR, "get-rows: VTCM reservation %zu is too small, needed %zu\n",
+             octx->ctx->vtcm_size, grctx.vtcm_layout.total_bytes);
+        return HTP_STATUS_INVAL_PARAMS;
+    }
 
     const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32);
 
     work_queue_func_t q_func = NULL;
-    if (kparams->use_dma) {
-        q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_st_int32_t : get_rows_thread_st_int64_t);
-    } else {
-        switch (octx->src[0]->type) {
-            case HTP_TYPE_F32:  q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_f32_int32_t  : get_rows_thread_f32_int64_t);  break;
-            case HTP_TYPE_F16:  q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_f16_int32_t  : get_rows_thread_f16_int64_t);  break;
-            case HTP_TYPE_Q8_0: q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_q8_0_int32_t : get_rows_thread_q8_0_int64_t); break;
-            case HTP_TYPE_I32:  q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_st_int32_t   : get_rows_thread_st_int64_t);   break;
-            default:            return HTP_STATUS_NO_SUPPORT;
-        }
+    switch (kparams->kernel_type) {
+        case HTP_GET_ROWS_KERNEL_SAMETYPE:
+            q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_st_int32_t : get_rows_thread_st_int64_t);
+            break;
+        case HTP_GET_ROWS_KERNEL_TILED:
+            q_func = get_rows_thread_tiled;
+            break;
+        case HTP_GET_ROWS_KERNEL_FLAT:
+            switch (octx->src[0]->type) {
+                case HTP_TYPE_F16:  q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_f16_int32_t  : get_rows_thread_f16_int64_t);  break;
+                case HTP_TYPE_Q8_0: q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_q8_0_int32_t : get_rows_thread_q8_0_int64_t); break;
+                default:            return HTP_STATUS_NO_SUPPORT;
+            }
+            break;
+        default:
+            return HTP_STATUS_NO_SUPPORT;
     }
 
-    FARF(HIGH, "get-rows: (%ux%ux%ux%u) x (%ux%ux%ux%u) -> (%ux%ux%ux%u) : src0-vtcm-size %zu dst-vtcm-size %zu use-dma %d n-threads %d\n",
+    FARF(HIGH, "get-rows: (%ux%ux%ux%u) x (%ux%ux%ux%u) -> (%ux%ux%ux%u) : src0-vtcm-size %zu dst-vtcm-size %zu kernel-type %d n-threads %d\n",
          octx->src[0]->ne[0], octx->src[0]->ne[1], octx->src[0]->ne[2], octx->src[0]->ne[3],
          octx->src[1]->ne[0], octx->src[1]->ne[1], octx->src[1]->ne[2], octx->src[1]->ne[3],
          octx->dst->ne[0], octx->dst->ne[1], octx->dst->ne[2], octx->dst->ne[3],
          grctx.vtcm_layout.src0_bytes_per_thread * n_threads,
          grctx.vtcm_layout.dst_bytes_per_thread  * n_threads,
-         kparams->use_dma, n_threads);
+         kparams->kernel_type, n_threads);
 
     work_queue_run(octx->ctx->work_queue, q_func, &grctx, n_threads);
     return HTP_STATUS_OK;

ggml/src/ggml-hexagon/htp/get-rows-ops.h

@@ -1,11 +1,21 @@
 #ifndef HTP_GET_ROWS_OPS_H
 #define HTP_GET_ROWS_OPS_H
 
+#include <stdbool.h>
+#include <string.h>
+
 #include "hex-fastdiv.h"
+#include "matmul-ops.h"
+
+enum htp_get_rows_kernel_type {
+    HTP_GET_ROWS_KERNEL_SAMETYPE = 0,
+    HTP_GET_ROWS_KERNEL_TILED,
+    HTP_GET_ROWS_KERNEL_FLAT,
+};
 
 struct htp_get_rows_kernel_params {
     int32_t  n_threads;
-    int32_t  use_dma;
+    int32_t  kernel_type;
     int32_t  chunks_per_row;
     int32_t  chunk_size;
     int32_t  total_tasks;
@@ -34,19 +44,37 @@ struct htp_get_rows_vtcm_layout {
 
 static inline void htp_get_rows_vtcm_layout_build(
     struct htp_get_rows_vtcm_layout * vtcm_layout,
+    int kernel_type,
     int type,
     uint32_t ne00,
     uint32_t n_threads) {
 
+    if (kernel_type == HTP_GET_ROWS_KERNEL_SAMETYPE) {
+        memset(vtcm_layout, 0, sizeof(*vtcm_layout));
+        return;
+    }
+
+    if (kernel_type == HTP_GET_ROWS_KERNEL_TILED) {
+        const size_t tile_size   = type == HTP_TYPE_Q4_0 ? HTP_MM_WEIGHT_TILE_SIZE_Q4_0 : HTP_MM_WEIGHT_TILE_SIZE_Q8_0;
+        const size_t tile_stride = (tile_size + 127) & ~127;
+        const uint32_t n_k_tiles = ne00 / HTP_MM_HMX_TILE_N_COLS;
+        const size_t row_tiles_size = n_k_tiles > 0 ? (n_k_tiles * tile_stride) : tile_stride;
+        vtcm_layout->src0_spad_half_size = (row_tiles_size + 255) & ~255;
+        vtcm_layout->dst_spad_half_size  = (ne00 * sizeof(float) + 255) & ~255;
+        vtcm_layout->src0_bytes_per_thread = 2 * vtcm_layout->src0_spad_half_size;
+        vtcm_layout->dst_bytes_per_thread  = 2 * vtcm_layout->dst_spad_half_size;
+        vtcm_layout->off_src0 = 0;
+        vtcm_layout->off_dst  = vtcm_layout->src0_bytes_per_thread * n_threads;
+        vtcm_layout->total_bytes = vtcm_layout->off_dst + vtcm_layout->dst_bytes_per_thread * n_threads;
+        return;
+    }
+
     uint32_t src0_row_size = 0;
     switch (type) {
-        case 0: // HTP_TYPE_F32
-            src0_row_size = ne00 * 4;
-            break;
-        case 1: // HTP_TYPE_F16
+        case HTP_TYPE_F16:
             src0_row_size = ne00 * 2;
             break;
-        case 8: // HTP_TYPE_Q8_0
+        case HTP_TYPE_Q8_0:
             src0_row_size = (ne00 / 32) * 34;
             break;
         default:

ggml/src/ggml-hexagon/htp/htp-ctx.h

@@ -150,7 +150,9 @@ int op_matmul_nx(struct htp_ops_context * octx);
 int op_matmul_id_nx(struct htp_ops_context * octx);
 int op_binary(struct htp_ops_context * octx);
 int op_unary(struct htp_ops_context * octx);
+int op_sum(struct htp_ops_context * octx);
 int op_sum_rows(struct htp_ops_context * octx);
+int op_argmax(struct htp_ops_context * octx);
 int op_activations(struct htp_ops_context * octx);
 int op_softmax(struct htp_ops_context * octx);
 int op_add_id(struct htp_ops_context * octx);

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

@@ -69,6 +69,7 @@ enum htp_op_code {
     HTP_OP_UNARY_ABS,
     HTP_OP_UNARY_LOG,
     HTP_OP_UNARY_RELU,
+    HTP_OP_UNARY_STEP,
     HTP_OP_GLU_SWIGLU,
     HTP_OP_GLU_SWIGLU_OAI,
     HTP_OP_GLU_GEGLU,
@@ -86,6 +87,7 @@ enum htp_op_code {
     HTP_OP_TOP_K,
     HTP_OP_SQR,
     HTP_OP_SQRT,
+    HTP_OP_SUM,
     HTP_OP_SUM_ROWS,
     HTP_OP_SSM_CONV,
     HTP_OP_REPEAT,
@@ -108,6 +110,7 @@ enum htp_op_code {
     HTP_OP_GLU_SWIGLU_CLAMP,
     HTP_OP_MDEV_GROUP,
     HTP_OP_ROLL,
+    HTP_OP_ARGMAX,
 
     HTP_OP_INVALID
 };

ggml/src/ggml-hexagon/htp/hvx-arith.h

@@ -579,6 +579,58 @@ static inline void hvx_abs_f16(uint8_t * restrict dst, const uint8_t * restrict
     }
 }
 
+//
+// Step
+//
+
+static inline void hvx_step_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
+    assert((unsigned long) dst % 128 == 0);
+    assert((unsigned long) src % 128 == 0);
+
+    HVX_Vector * restrict vdst = (HVX_Vector *) dst;
+    HVX_Vector * restrict vsrc = (HVX_Vector *) src;
+
+    const uint32_t elem_size = sizeof(float);
+    const uint32_t epv       = 128 / elem_size;
+    const uint32_t nvec      = n / epv;
+    const uint32_t nloe      = n % epv;
+
+    uint32_t i = 0;
+
+    _Pragma("unroll(4)")
+    for (; i < nvec; i++) {
+        vdst[i] = hvx_vec_step_f32(vsrc[i]);
+    }
+    if (nloe) {
+        HVX_Vector v = hvx_vec_step_f32(vsrc[i]);
+        hvx_vec_store_a((void *) &vdst[i], nloe * elem_size, v);
+    }
+}
+
+static inline void hvx_step_f16_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) {
+    assert((unsigned long) dst % 128 == 0);
+    assert((unsigned long) src % 128 == 0);
+
+    HVX_Vector * restrict vdst = (HVX_Vector *) dst;
+    HVX_Vector * restrict vsrc = (HVX_Vector *) src;
+
+    const uint32_t elem_size = sizeof(_Float16);
+    const uint32_t epv       = 128 / elem_size;
+    const uint32_t nvec      = n / epv;
+    const uint32_t nloe      = n % epv;
+
+    uint32_t i = 0;
+
+    _Pragma("unroll(4)")
+    for (; i < nvec; i++) {
+        vdst[i] = hvx_vec_step_f16(vsrc[i]);
+    }
+    if (nloe) {
+        HVX_Vector v = hvx_vec_step_f16(vsrc[i]);
+        hvx_vec_store_a((void *) &vdst[i], nloe * elem_size, v);
+    }
+}
+
 //
 // Square
 //

ggml/src/ggml-hexagon/htp/hvx-base.h

@@ -111,6 +111,20 @@ static inline HVX_Vector hvx_vec_neg_f32(HVX_Vector v) {
 #endif  // __HVX_ARCH__ > 75
 }
 
+static inline HVX_Vector hvx_vec_step_f32(HVX_Vector v) {
+    const HVX_Vector zero = Q6_V_vzero();
+    const HVX_Vector one  = hvx_vec_splat_f32(1.0f);
+    HVX_VectorPred q = Q6_Q_vcmp_gt_VsfVsf(v, zero);
+    return Q6_V_vmux_QVV(q, one, zero);
+}
+
+static inline HVX_Vector hvx_vec_step_f16(HVX_Vector v) {
+    const HVX_Vector zero = Q6_V_vzero();
+    const HVX_Vector one  = hvx_vec_splat_f16((_Float16) 1.0f);
+    HVX_VectorPred q = Q6_Q_vcmp_gt_VhfVhf(v, zero);
+    return Q6_V_vmux_QVV(q, one, zero);
+}
+
 static inline HVX_VectorPred hvx_vec_is_nan_f16(HVX_Vector v) {
     const HVX_Vector vnan_exp  = Q6_Vh_vsplat_R(0x7C00);
     const HVX_Vector vnan_frac = Q6_Vh_vsplat_R(0x7FFF);

ggml/src/ggml-hexagon/htp/hvx-reduce.h

@@ -326,6 +326,79 @@ static inline int32_t hvx_reduce_max_i32(const uint8_t * restrict src, const int
     }
 }
 
+static inline void hvx_argmax_f32(
+    const float * restrict src,
+    uint32_t n,
+    uint32_t offset,
+    float * out_val,
+    int32_t * out_idx
+) {
+    if (n == 0) {
+        *out_val = -INFINITY;
+        *out_idx = (int32_t) offset;
+        return;
+    }
+
+    if (n < 32 || !hex_is_aligned((void *) src, 128)) {
+        float best_val = src[0];
+        int32_t best_idx = (int32_t) offset;
+        for (uint32_t i = 1; i < n; i++) {
+            if (src[i] > best_val) {
+                best_val = src[i];
+                best_idx = (int32_t) (offset + i);
+            }
+        }
+        *out_val = best_val;
+        *out_idx = best_idx;
+        return;
+    }
+
+    static const int32_t c_lane_idx[32] __attribute__((aligned(128))) = {
+         0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15,
+        16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31
+    };
+
+    const HVX_Vector v_init_idx = *(const HVX_Vector *) c_lane_idx;
+    const HVX_Vector v_step     = Q6_V_vsplat_R(32);
+    HVX_Vector v_cur_idx        = Q6_Vw_vadd_VwVw(v_init_idx, Q6_V_vsplat_R((int32_t) offset));
+    HVX_Vector v_max_val        = hvx_vec_splat_f32(-INFINITY);
+    HVX_Vector v_max_idx        = v_cur_idx;
+
+    const HVX_Vector * vsrc = (const HVX_Vector *) src;
+    const uint32_t nvec = n / 32;
+
+    for (uint32_t vi = 0; vi < nvec; vi++) {
+        HVX_Vector v = vsrc[vi];
+        HVX_VectorPred pred = Q6_Q_vcmp_gt_VsfVsf(v, v_max_val);
+        v_max_val = Q6_V_vmux_QVV(pred, v, v_max_val);
+        v_max_idx = Q6_V_vmux_QVV(pred, v_cur_idx, v_max_idx);
+        v_cur_idx = Q6_Vw_vadd_VwVw(v_cur_idx, v_step);
+    }
+
+    HVX_VectorAlias u_val, u_idx;
+    u_val.v = v_max_val;
+    u_idx.v = v_max_idx;
+
+    float best_val = u_val.fp32[0];
+    int32_t best_idx = (int32_t) u_idx.w[0];
+    for (int i = 1; i < 32; i++) {
+        if (u_val.fp32[i] > best_val) {
+            best_val = u_val.fp32[i];
+            best_idx = (int32_t) u_idx.w[i];
+        }
+    }
+
+    for (uint32_t i = nvec * 32; i < n; i++) {
+        if (src[i] > best_val) {
+            best_val = src[i];
+            best_idx = (int32_t) (offset + i);
+        }
+    }
+
+    *out_val = best_val;
+    *out_idx = best_idx;
+}
+
 #undef hvx_reduce_loop_body
 #undef HVX_REDUCE_MAX_OP
 #undef HVX_REDUCE_SUM_OP

ggml/src/ggml-hexagon/htp/main.c

@@ -857,6 +857,7 @@ static int execute_op(struct htp_ops_context * octx) {
         case HTP_OP_UNARY_ABS:
         case HTP_OP_UNARY_LOG:
         case HTP_OP_UNARY_RELU:
+        case HTP_OP_UNARY_STEP:
         case HTP_OP_L2_NORM:
             return op_unary(octx);
 
@@ -882,6 +883,9 @@ static int execute_op(struct htp_ops_context * octx) {
         case HTP_OP_GET_ROWS:
             return op_get_rows(octx);
 
+        case HTP_OP_SUM:
+            return op_sum(octx);
+
         case HTP_OP_SUM_ROWS:
             return op_sum_rows(octx);
 
@@ -898,6 +902,9 @@ static int execute_op(struct htp_ops_context * octx) {
         case HTP_OP_TOP_K:
             return op_top_k(octx);
 
+        case HTP_OP_ARGMAX:
+            return op_argmax(octx);
+
         case HTP_OP_SSM_CONV:
             return op_ssm_conv(octx);
 

ggml/src/ggml-hexagon/htp/sum-rows-ops.c

@@ -83,16 +83,17 @@ static void sum_rows_thread_f32(unsigned int nth, unsigned int ith, void *data)
     htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start_row);
 
     for (uint32_t ir = 0; ir < n_rows; ir++) {
-        const float * restrict src_local = src_th + (ir * (src_stride / sizeof(float)));
+        const float * restrict src_local = (const float *) ((const uint8_t *) src_th + ir * src_stride);
+        float       * restrict dst_local = (float *)       ((uint8_t *)       dst_th + ir * dst_stride);
 
         if (ir + 1 < n_rows) {
-            hex_l2fetch(src_local + (src_stride / sizeof(float)), src_stride, src_stride, 1);
+            hex_l2fetch((const uint8_t *) src_local + src_stride, src_stride, src_stride, 1);
         }
 
         if (opt_path) {
-            dst_th[ir] = hvx_reduce_sum_f32_a((const uint8_t *) src_local, ne00);
+            *dst_local = hvx_reduce_sum_f32_a((const uint8_t *) src_local, ne00);
         } else {
-            dst_th[ir] = hvx_reduce_sum_f32((const uint8_t *) src_local, ne00);
+            *dst_local = hvx_reduce_sum_f32((const uint8_t *) src_local, ne00);
         }
     }
 
@@ -152,3 +153,258 @@ int op_sum_rows(struct htp_ops_context * octx) {
 
     return HTP_STATUS_OK;
 }
+
+struct sum_context {
+    struct htp_ops_context * octx;
+    const float            * src_data;
+    float                    partial_sums[HTP_MAX_NTHREADS];
+    uint32_t                 total_elems;
+    uint32_t                 elems_per_thread;
+};
+
+static void sum_thread_f32(unsigned int nth, unsigned int ith, void * data) {
+    struct sum_context * sctx = (struct sum_context *) data;
+    const uint32_t start = sctx->elems_per_thread * ith;
+    const uint32_t end   = MIN(start + sctx->elems_per_thread, sctx->total_elems);
+
+    if (start >= end) {
+        sctx->partial_sums[ith] = 0.0f;
+        return;
+    }
+
+    const uint32_t n = end - start;
+    const float * src = sctx->src_data + start;
+
+    struct htp_thread_trace * tr = &sctx->octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start);
+
+    hex_l2fetch_block((const void *) src, n * sizeof(float));
+
+    sctx->partial_sums[ith] = hvx_reduce_sum_f32((const uint8_t *) src, n);
+
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start);
+}
+
+int op_sum(struct htp_ops_context * octx) {
+    const struct htp_tensor * src0 = octx->src[0];
+    const struct htp_tensor * dst  = octx->dst;
+
+    if (src0->type != HTP_TYPE_F32) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
+    if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
+    if (octx->ctx->mdev.count > 1 && octx->ctx->mdev.idx > 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t total_elems = (uint32_t) (src0->ne[0] * src0->ne[1] * src0->ne[2] * src0->ne[3]);
+    if (total_elems == 0) {
+        ((float *) dst->data)[0] = 0.0f;
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = (total_elems >= 1024) ? MIN(octx->n_threads, HTP_MAX_NTHREADS) : 1;
+    const uint32_t raw_chunk = (total_elems + n_threads - 1) / n_threads;
+    const uint32_t elems_per_thread = hex_round_up(raw_chunk, 32);
+
+    struct sum_context sctx = {
+        .octx             = octx,
+        .src_data         = (const float *) src0->data,
+        .total_elems      = total_elems,
+        .elems_per_thread = elems_per_thread,
+    };
+
+    work_queue_run(octx->ctx->work_queue, sum_thread_f32, &sctx, n_threads);
+
+    float sum = 0.0f;
+    for (uint32_t i = 0; i < n_threads; i++) {
+        sum += sctx.partial_sums[i];
+    }
+    ((float *) dst->data)[0] = sum;
+
+    return HTP_STATUS_OK;
+}
+
+static inline void argmax_slice_f32(
+    const float * restrict src,
+    uint32_t n,
+    uint32_t offset,
+    float * out_val,
+    int32_t * out_idx
+) {
+    hvx_argmax_f32(src, n, offset, out_val, out_idx);
+}
+
+struct argmax_context {
+    struct htp_ops_context * octx;
+    const float            * src_data;
+    int32_t                * dst_data;
+    uint32_t                 ne00;
+    uint32_t                 src_stride;
+    uint32_t                 dst_stride;
+    uint32_t                 row_start;
+    uint32_t                 nrows;
+    uint32_t                 rows_per_thread;
+    uint32_t                 elems_per_thread;
+
+    float                    partial_max[HTP_MAX_NTHREADS];
+    int32_t                  partial_idx[HTP_MAX_NTHREADS];
+};
+
+static void argmax_thread_single_row(unsigned int nth, unsigned int ith, void * data) {
+    struct argmax_context * actx = (struct argmax_context *) data;
+    const uint32_t start = actx->elems_per_thread * ith;
+    const uint32_t end   = MIN(start + actx->elems_per_thread, actx->ne00);
+
+    if (start >= end) {
+        actx->partial_max[ith] = -INFINITY;
+        actx->partial_idx[ith] = 0;
+        return;
+    }
+
+    const uint32_t n = end - start;
+    const float * src = actx->src_data + start;
+
+    struct htp_thread_trace * tr = &actx->octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start);
+
+    hex_l2fetch_block((const void *) src, n * sizeof(float));
+
+    float max_val;
+    int32_t max_idx;
+    argmax_slice_f32(src, n, start, &max_val, &max_idx);
+
+    actx->partial_max[ith] = max_val;
+    actx->partial_idx[ith] = max_idx;
+
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) start);
+}
+
+static void argmax_thread_multi_row(unsigned int nth, unsigned int ith, void * data) {
+    struct argmax_context * actx = (struct argmax_context *) data;
+    const uint32_t r0 = actx->row_start + actx->rows_per_thread * ith;
+    const uint32_t r1 = MIN(r0 + actx->rows_per_thread, actx->row_start + actx->nrows);
+
+    if (r0 >= r1) {
+        return;
+    }
+
+    struct htp_thread_trace * tr = &actx->octx->ctx->trace[ith];
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r0);
+
+    for (uint32_t r = r0; r < r1; r++) {
+        const float * src_row = (const float *) ((const uint8_t *) actx->src_data + r * actx->src_stride);
+        int32_t * dst_val     = (int32_t *) ((uint8_t *) actx->dst_data + r * actx->dst_stride);
+
+        hex_l2fetch_block((const void *) src_row, actx->ne00 * sizeof(float));
+
+        float max_val;
+        int32_t max_idx;
+        argmax_slice_f32(src_row, actx->ne00, 0, &max_val, &max_idx);
+
+        *dst_val = max_idx;
+    }
+
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r0);
+}
+
+int op_argmax(struct htp_ops_context * octx) {
+    const struct htp_tensor * src0 = octx->src[0];
+    const struct htp_tensor * dst  = octx->dst;
+
+    if (src0->type != HTP_TYPE_F32 || dst->type != HTP_TYPE_I32) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
+    if (htp_tensor_is_extended(src0) || htp_tensor_is_extended(dst)) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
+    const uint32_t ne00 = src0->ne[0];
+    const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
+
+    if (ne00 == 0 || src0_nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    if (src0_nrows == 1) {
+        if (octx->ctx->mdev.count > 1 && octx->ctx->mdev.idx > 0) {
+            return HTP_STATUS_OK;
+        }
+
+        if (ne00 == 1) {
+            ((int32_t *) dst->data)[0] = 0;
+            return HTP_STATUS_OK;
+        }
+
+        const uint32_t n_threads = (ne00 >= 1024) ? MIN(octx->n_threads, HTP_MAX_NTHREADS) : 1;
+        const uint32_t raw_chunk = (ne00 + n_threads - 1) / n_threads;
+        const uint32_t elems_per_thread = hex_round_up(raw_chunk, 32);
+
+        struct argmax_context actx = {
+            .octx             = octx,
+            .src_data         = (const float *) src0->data,
+            .dst_data         = (int32_t *) dst->data,
+            .ne00             = ne00,
+            .src_stride       = src0->nb[1] > 0 ? (uint32_t) src0->nb[1] : (uint32_t) (ne00 * sizeof(float)),
+            .dst_stride       = dst->nb[0] > 0 ? (uint32_t) dst->nb[0] : (uint32_t) sizeof(int32_t),
+            .row_start        = 0,
+            .nrows            = 1,
+            .rows_per_thread  = 1,
+            .elems_per_thread = elems_per_thread,
+        };
+
+        work_queue_run(octx->ctx->work_queue, argmax_thread_single_row, &actx, n_threads);
+
+        float best_val = actx.partial_max[0];
+        int32_t best_idx = actx.partial_idx[0];
+        for (uint32_t i = 1; i < n_threads; i++) {
+            if (actx.partial_max[i] > best_val) {
+                best_val = actx.partial_max[i];
+                best_idx = actx.partial_idx[i];
+            }
+        }
+        ((int32_t *) dst->data)[0] = best_idx;
+        return HTP_STATUS_OK;
+    }
+
+    uint32_t row_start = 0;
+    uint32_t nrows     = src0_nrows;
+
+    if (octx->ctx->mdev.count > 1) {
+        const bool can_split = htp_tensor_mdev_data_aligned(dst) &&
+                               htp_tensor_is_contiguous(dst, sizeof(int32_t));
+        const uint32_t elems_per_chunk = can_split ? 32 : 0;
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(
+            src0_nrows, elems_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        row_start = range.start;
+        nrows     = range.count;
+    }
+
+    if (nrows == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = MIN(octx->n_threads, nrows);
+    const uint32_t rows_per_thread = (nrows + n_threads - 1) / n_threads;
+
+    struct argmax_context actx = {
+        .octx             = octx,
+        .src_data         = (const float *) src0->data,
+        .dst_data         = (int32_t *) dst->data,
+        .ne00             = ne00,
+        .src_stride       = src0->nb[1] > 0 ? (uint32_t) src0->nb[1] : (uint32_t) (ne00 * sizeof(float)),
+        .dst_stride       = dst->nb[0] > 0 ? (uint32_t) dst->nb[0] : (uint32_t) sizeof(int32_t),
+        .row_start        = row_start,
+        .nrows            = nrows,
+        .rows_per_thread  = rows_per_thread,
+        .elems_per_thread = 0,
+    };
+
+    work_queue_run(octx->ctx->work_queue, argmax_thread_multi_row, &actx, n_threads);
+    return HTP_STATUS_OK;
+}

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

@@ -409,6 +409,20 @@ static void log_f16(const void * restrict src,
     }
 }
 
+static void step_f16(const void * restrict src,
+                     void * restrict dst,
+                     const uint32_t num_rows,
+                     const struct htp_unary_context * uctx) {
+    htp_unary_op_preamble;
+
+    for (uint32_t ir = 0; ir < num_rows; ir++) {
+        const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+        uint8_t * restrict dst_local       = (uint8_t *)dst + (ir * dst_row_size_aligned);
+
+        hvx_step_f16_aa((uint8_t *) dst_local, (const uint8_t *) src_local, ne0);
+    }
+}
+
 static void l2_norm_f16(const void * restrict src,
                         void * restrict dst,
                         const uint32_t num_rows,
@@ -662,6 +676,20 @@ static void relu_f32(const void * restrict src,
     }
 }
 
+static void step_f32(const void * restrict src,
+                     void * restrict dst,
+                     const uint32_t num_rows,
+                     const struct htp_unary_context * uctx) {
+    htp_unary_op_preamble;
+
+    for (uint32_t ir = 0; ir < num_rows; ir++) {
+        const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+        uint8_t * restrict dst_local       = (uint8_t *)dst + (ir * dst_row_size_aligned);
+
+        hvx_step_f32_aa(dst_local, src_local, ne0);
+    }
+}
+
 static void log_f32(const void * restrict src,
                     void * restrict dst,
                     const uint32_t num_rows,
@@ -767,6 +795,11 @@ static void tile_relu_f32(void * restrict dst, const void * restrict src, uint32
     hvx_max_scalar_f32((uint8_t *) dst, (const uint8_t *) src, 0.0f, tw);
 }
 
+static void tile_step_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) {
+    (void) uctx;
+    hvx_step_f32_aa((uint8_t *) dst, (const uint8_t *) src, tw);
+}
+
 static void tri_apply_tile_f32(const void * restrict src, void * restrict dst,
                                uint32_t tile_elems, uint32_t col_start, uint32_t i01,
                                uint32_t ne0, int32_t ttype) {
@@ -1486,6 +1519,7 @@ static int execute_op_unary(struct htp_ops_context * octx) {
         case HTP_OP_UNARY_ABS:       op_type = is_f16 ? "abs-f16"      : "abs-f32";          break;
         case HTP_OP_UNARY_LOG:       op_type = is_f16 ? "log-f16"      : "log-f32";          break;
         case HTP_OP_UNARY_RELU:      op_type = "relu-f32";                                   break;
+        case HTP_OP_UNARY_STEP:      op_type = is_f16 ? "step-f16"     : "step-f32";         break;
         case HTP_OP_L2_NORM:         op_type = is_f16 ? "l2norm-f16"   : "l2norm-f32";       break;
         case HTP_OP_TRI:             op_type = "tri-f32";                                    break;
         default:
@@ -1506,6 +1540,7 @@ static int execute_op_unary(struct htp_ops_context * octx) {
             case HTP_OP_L2_NORM:
             case HTP_OP_UNARY_ABS:
             case HTP_OP_UNARY_LOG:
+            case HTP_OP_UNARY_STEP:
                 break;
             default:
                 FARF(ERROR, "unary-%s: not supported for F16\n", op_type);
@@ -1634,6 +1669,7 @@ static int execute_op_unary(struct htp_ops_context * octx) {
             case HTP_OP_UNARY_ABS:       compute_func = (void *) tile_abs_f32;            break;
             case HTP_OP_UNARY_LOG:       compute_func = (void *) tile_log_f32;            break;
             case HTP_OP_UNARY_RELU:      compute_func = (void *) tile_relu_f32;           break;
+            case HTP_OP_UNARY_STEP:      compute_func = (void *) tile_step_f32;           break;
             case HTP_OP_TRI:
                 task_func    = unary_thread_tiled_tri_f32;
                 compute_func = (void *) tri_apply_tile_f32;
@@ -1652,6 +1688,7 @@ static int execute_op_unary(struct htp_ops_context * octx) {
             case HTP_OP_L2_NORM:         compute_func = (void *) l2_norm_f16;             break;
             case HTP_OP_UNARY_ABS:       compute_func = (void *) abs_f16;                 break;
             case HTP_OP_UNARY_LOG:       compute_func = (void *) log_f16;                 break;
+            case HTP_OP_UNARY_STEP:      compute_func = (void *) step_f16;                break;
             default:                     break;
         }
     } else {
@@ -1678,6 +1715,7 @@ static int execute_op_unary(struct htp_ops_context * octx) {
             case HTP_OP_UNARY_ABS:       compute_func = (void *) abs_f32;                 break;
             case HTP_OP_UNARY_LOG:       compute_func = (void *) log_f32;                 break;
             case HTP_OP_UNARY_RELU:      compute_func = (void *) relu_f32;                break;
+            case HTP_OP_UNARY_STEP:      compute_func = (void *) step_f32;                break;
             case HTP_OP_L2_NORM:         compute_func = (void *) l2_norm_f32;             break;
             case HTP_OP_TRI:
                 task_func    = unary_thread_tri_f32;

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

@@ -59,6 +59,7 @@ static inline bool htp_op_is_unary(uint32_t opcode) {
         case HTP_OP_UNARY_ABS:
         case HTP_OP_UNARY_LOG:
         case HTP_OP_UNARY_RELU:
+        case HTP_OP_UNARY_STEP:
         case HTP_OP_L2_NORM:
         case HTP_OP_TRI:
             return true;

src/llama-arch.cpp

@@ -722,7 +722,7 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
     {LLM_TENSOR_TOKEN_EMBD,                 {LLM_TENSOR_LAYER_INPUT,     GGML_OP_GET_ROWS}},
     {LLM_TENSOR_POS_EMBD,                   {LLM_TENSOR_LAYER_INPUT,     GGML_OP_GET_ROWS}},
     {LLM_TENSOR_TOKEN_TYPES,                {LLM_TENSOR_LAYER_INPUT,     GGML_OP_GET_ROWS}},
-    {LLM_TENSOR_HRM_Z_L_INIT,               {LLM_TENSOR_LAYER_INPUT,     GGML_OP_ADD}},
+    {LLM_TENSOR_HRM_Z_L_INIT,               {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}},
     {LLM_TENSOR_TOKEN_EMBD_NORM,            {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},  // do the norms on the first layer (not the input layer)
     {LLM_TENSOR_OUTPUT,                     {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_MUL_MAT}},
     {LLM_TENSOR_CLS,                        {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_MUL_MAT}},

src/models/hrm-text.cpp

@@ -43,7 +43,7 @@ void llama_model_hrm_text::load_arch_tensors(llama_model_loader &) {
         output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED);
     }
 
-    hrm_z_l_init = create_tensor(tn(LLM_TENSOR_HRM_Z_L_INIT), { n_embd }, 0);
+    hrm_z_l_init = create_tensor(tn(LLM_TENSOR_HRM_Z_L_INIT, 0), { n_embd }, 0);
 
     const int lps = hparams.n_hrm_layers_per_stack;
 

tests/test-backend-ops.cpp

@@ -9830,6 +9830,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         add_test_bin_bcast(type, {5120, 1, 1, 1}, {1, 256, 1, 1});
         add_test_bin_bcast(type, {640, 1, 1, 1}, {1, 1, 1, 1});
         add_test_bin_bcast(type, {64, 262144, 1, 1}, {1, 1, 1, 1});
+        add_test_bin_bcast(type, {128, 1, 8, 1}, {1, 4, 1, 1});
         //add_test_bin_bcast(type, {3, 3, 2560, 1280}, {1, 1, 1, 1});
         //add_test_bin_bcast(type, {3, 3, 2560, 1280}, {2, 1, 1, 1});
     }
@@ -9979,6 +9980,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     }
 
     test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 16, 1, 1024, 1, 32, 4)); // Mamba-1
+    test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 96, 64, 128, 8, 1, 1)); // Nemotron-3-Puzzle decode (scan path, unused warp in last block)
+    test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 96, 64, 16, 8, 300, 2, false, /*K=*/1, /*weak_decay=*/true)); // d_state=96 SSD multi-chunk (partial 2nd chunk, 2 seqs)
+    test_cases.emplace_back(new test_ssm_scan_rollback(GGML_TYPE_F32, 96, 64, 16, 8, 8, 2, /*K=*/3)); // d_state=96 rollback snapshots match prefix states
     test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 16, 2, 32, 4)); // Mamba-2
     test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 256, 64,  8, 2, 32, 4)); // Falcon-H1
     test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 128, 4, 4, 16, 2, true)); // x/B/C overlap
@@ -10017,6 +10021,12 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 128, 32, 128));
     test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 128, 4, 128, {2, 3}));
     test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 256, 512, 256)); // many rows
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 64, 1, 64));
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 128, 1, 128));
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 256, 1, 256));
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 512, 1, 512));
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 128, 32, 128));
+    test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F16, 256, 512, 256)); // many rows
     test_cases.emplace_back(new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 32, 1, 32)); // too small (N<64)
     test_cases.emplace_back(
         new test_mul_mat_hadamard(GGML_TYPE_F32, GGML_TYPE_F32, 16384, 1, 16384));               // too big (N>8192)

tools/llama-bench/llama-bench.cpp

@@ -1097,7 +1097,7 @@ static cmd_params parse_cmd_params(int argc, char ** argv) {
             p.hf_token      = params.hf_token;
             p.offline       = params.offline;
             p.model.hf_repo = params.hf_repo[i];
-            if (!params.hf_file.empty() && !params.hf_file[i].empty()) {
+            if (i < params.hf_file.size() && !params.hf_file[i].empty()) {
                 p.model.hf_file = params.hf_file[i];
             }