mimo2 : expose t_layer_inp for the DFlash drafter

Assisted-by: pi:llama.cpp/Qwen3.8-27B
convert : handle dict/list-wrapped mask_embedding.pt in DFlash drafter
2026-10-27 06:57:35 -05:00 · 2026-09-28 18:05:48 +03:00 · 2026-09-28 16:39:22 +03:00 · 2026-09-28 16:32:03 +03:00
281 changed files with 12478 additions and 28012 deletions
@@ -1,4 +1,4 @@
-ARG ONEAPI_VERSION=2026.1.1-devel-ubuntu24.04
+ARG ONEAPI_VERSION=2025.3.3-0-devel-ubuntu24.04
 ARG BUILD_DATE=N/A
 ARG APP_VERSION=N/A
 ARG APP_REVISION=N/A
@@ -19,7 +19,7 @@ RUN npm ci
 COPY tools/ui/ ./
 RUN LLAMA_BUILD_NUMBER="$APP_VERSION" npm run build

-FROM docker.io/intel/oneapi-toolkit:$ONEAPI_VERSION AS build
+FROM docker.io/intel/deep-learning-essentials:$ONEAPI_VERSION AS build

 ARG GGML_SYCL_F16=ON
 ARG LEVEL_ZERO_VERSION=1.28.2
@@ -59,7 +59,7 @@ RUN mkdir -p /app/full \
    && cp requirements.txt /app/full \
    && cp .devops/tools.sh /app/full/tools.sh

-FROM docker.io/intel/oneapi-toolkit:$ONEAPI_VERSION AS base
+FROM docker.io/intel/deep-learning-essentials:$ONEAPI_VERSION AS base

 ARG BUILD_DATE=N/A
 ARG APP_VERSION=N/A
@@ -1,9 +1,9 @@
 ARG UBUNTU_VERSION=22.04
 # This needs to generally match the container host's environment.
 # Target the MUSA build image
-ARG BASE_MUSA_DEV_CONTAINER=registry.mthreads.com/mcconline/musa_sdk:5.2.0-devel-ubuntu${UBUNTU_VERSION}-s5000
+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=registry.mthreads.com/mcconline/musa_sdk:5.2.0-runtime-ubuntu${UBUNTU_VERSION}-s5000
+ARG BASE_MUSA_RUN_CONTAINER=${BASE_MUSA_DEV_CONTAINER}

 ARG BUILD_DATE=N/A
 ARG APP_VERSION=N/A
@@ -36,10 +36,7 @@ RUN apt-get update && \
    python3-pip \
    git \
    libssl-dev \
-    libgomp1 \
-    musa-mualg-5-2 \
-    musa-muthrust-5-2 \
-    libmthreads-compute
+    libgomp1

 WORKDIR /app

@@ -82,16 +79,13 @@ LABEL org.opencontainers.image.created=$BUILD_DATE \
      org.opencontainers.image.source=$IMAGE_SOURCE

 RUN apt-get update \
-    && apt-get install -y libgomp1 curl ffmpeg libmthreads-compute \
+    && apt-get install -y libgomp1 curl ffmpeg \
    && apt autoremove -y \
    && apt clean -y \
    && rm -rf /tmp/* /var/tmp/* \
    && find /var/cache/apt/archives /var/lib/apt/lists -not -name lock -type f -delete \
    && find /var/cache -type f -delete

-# The MUSA runtime image does not register its library directory
-RUN echo "/usr/local/musa/lib" > /etc/ld.so.conf.d/musa-runtime.conf && ldconfig
-
 COPY --from=build /app/lib/ /app

 ### Full
@@ -68,7 +68,7 @@ jobs:
          hf_bucket: ggml-org/cache

      - name: Build with CMake
-        # TODO: Drop GGML_CUDA_CCCL_VERSION when this job uses CTK >= 13.5, which bundles CCCL >= 3.5.
+        # TODO: Remove GGML_CUDA_CUB_3DOT2 flag once CCCL 3.2 is bundled within CTK and that CTK version is used in this project
        run: |
          cmake -S . -B build -G Ninja \
            -DLLAMA_FATAL_WARNINGS=ON \
@@ -77,7 +77,7 @@ jobs:
            -DCMAKE_EXE_LINKER_FLAGS=-Wl,--allow-shlib-undefined \
            -DGGML_NATIVE=OFF \
            -DGGML_CUDA=ON \
-            -DGGML_CUDA_CCCL_VERSION=v3.4.3
+            -DGGML_CUDA_CUB_3DOT2=ON
          cmake --build build

      - name: ccache-buckets-save
@@ -145,7 +145,7 @@ jobs:

  musa:
    runs-on: ubuntu-22.04
-    container: registry.mthreads.com/mcconline/musa_sdk:5.2.0-devel-ubuntu22.04-s5000
+    container: registry.mthreads.com/mcconline/inference/pytorch:2.9.1.post1-py3.10-musa5.2.0-mp31-devel-ubuntu22.04-amd64

    steps:
      - name: Clone
@@ -156,7 +156,7 @@ jobs:
        id: depends
        run: |
          apt-get update
-          apt-get install -y build-essential git cmake libssl-dev jq python3-venv musa-mualg-5-2 musa-muthrust-5-2 libmthreads-compute
+          apt-get install -y build-essential git cmake libssl-dev jq python3-venv

      - name: ccache
        uses: ggml-org/ccache-action@v1.2.24
@@ -31,16 +31,15 @@ jobs:
    strategy:
      matrix:
        include:
-          # CTK >= 13.5 bundles CCCL >= 3.5; omit GGML_CUDA_CCCL_VERSION for those versions.
          - cuda: '12.4'
            arch: x64
-            defines: '-DGGML_CUDA_CCCL_VERSION=v3.4.3'
+            defines: '-DGGML_CUDA_CUB_3DOT2=ON'
          - cuda: '13.4'
            arch: x64
-            defines: '-DGGML_CUDA_CCCL_VERSION=v3.4.3'
+            defines: ''
          - cuda: '13.4'
            arch: arm64
-            defines: '-DGGML_CUDA_CCCL_VERSION=v3.4.3 -DCMAKE_TOOLCHAIN_FILE=cmake/arm64-windows-msvc-cuda.cmake'
+            defines: '-DCMAKE_TOOLCHAIN_FILE=cmake/arm64-windows-msvc-cuda.cmake'

    steps:
      - name: Clone
@@ -100,46 +100,6 @@ jobs:
          wget https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K-be.gguf
          ./bin/llama-completion -m stories260K-be.gguf -p "One day, Lily met a Shoggoth" -n 500 -c 256

-  ubuntu-26-zdnn-s390x:
-    name: ubuntu-26-zdnn-s390x
-    runs-on: ubuntu-24.04-s390x
-    container: ubuntu:26.04  # required to get GCC 15.1 and binutils 2.44
-    defaults:
-      run:
-        shell: bash
-
-    steps:
-      - name: Build Dependencies
-        id: build_depends
-        run: |
-          apt-get update
-          apt-get install -y --no-install-recommends \
-            build-essential cmake git ca-certificates \
-            libssl-dev libzdnn-dev
-
-      - name: Clone
-        id: checkout
-        uses: actions/checkout@v6
-
-      - name: Toolchain workaround (GCC 15)
-        run: |
-          apt-get install -y gcc-15 g++-15
-          echo "CC=gcc-15" >> "$GITHUB_ENV"
-          echo "CXX=g++-15" >> "$GITHUB_ENV"
-
-      - name: Build with zDNN Backend
-        id: cmake_build
-        run: |
-          cmake -B build \
-            -DLLAMA_FATAL_WARNINGS=ON \
-            -DGGML_NATIVE=OFF \
-            -DGGML_VXE=ON \
-            -DGGML_ZDNN=ON \
-            -DGGML_RPC=ON \
-            -DCMAKE_C_FLAGS="-march=arch15" \
-            -DCMAKE_CXX_FLAGS="-march=arch15"
-          time cmake --build build --config Release -j $(nproc)
-
  ubuntu-24-ppc64le:
    runs-on: ubuntu-24.04-ppc64le

@@ -48,7 +48,7 @@ jobs:

    env:
      ONEAPI_ROOT: /opt/intel/oneapi/
-      ONEAPI_INSTALLER_VERSION: "2026.1"
+      ONEAPI_INSTALLER_VERSION: "2025.3.3"
      LEVEL_ZERO_VERSION: "1.33.1"
      LEVEL_ZERO_UBUNTU_VERSION: "u24.04"

@@ -63,8 +63,8 @@ jobs:
        shell: bash
        run: |
          cd /tmp
-          wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/5996e26b-f48a-42b1-8db0-b002ad0bd8d7/intel-oneapi-toolkit-2026.1.1.33_offline.sh -O intel-oneapi-toolkit_offline.sh
-          sudo bash intel-oneapi-toolkit_offline.sh -s -a --silent --eula accept
+          wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/56f7923a-adb8-43f3-8b02-2b60fcac8cab/intel-deep-learning-essentials-2025.3.3.16_offline.sh -O intel-deep-learning-essentials_offline.sh
+          sudo bash intel-deep-learning-essentials_offline.sh -s -a --silent --eula accept

      - name: Install Level Zero SDK
        shell: bash
@@ -129,11 +129,11 @@ jobs:
        shell: bash

    env:
-      WINDOWS_BASEKIT_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/0cb67a0d-67f6-410b-868b-f4a0a17ff0cf/intel-oneapi-toolkit-2026.1.1.32_offline.exe
+      WINDOWS_BASEKIT_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/b60765d1-2b85-4e85-86b6-cb0e9563a699/intel-deep-learning-essentials-2025.3.3.18_offline.exe
      WINDOWS_DPCPP_MKL: intel.oneapi.win.cpp-dpcpp-common:intel.oneapi.win.mkl.devel:intel.oneapi.win.dnnl:intel.oneapi.win.tbb.devel
      LEVEL_ZERO_SDK_URL: https://github.com/oneapi-src/level-zero/releases/download/v1.33.1/level-zero-win-sdk-1.33.1.zip
      ONEAPI_ROOT: "C:/Program Files (x86)/Intel/oneAPI"
-      ONEAPI_INSTALLER_VERSION: "2026.1"
+      ONEAPI_INSTALLER_VERSION: "2025.3.3"
    steps:
      - name: Clone
        id: checkout
@@ -0,0 +1,71 @@
+name: Fusion
+
+on:
+  workflow_dispatch: # allows manual triggering
+  push:
+    branches:
+      - master
+    paths: [
+      '.github/workflows/fusion.yml',
+      'ggml/**',
+      'tests/fusion/**',
+      'tests/test-fusion.cpp',
+      'tests/test-llama-archs.cpp',
+      'src/models/**'
+    ]
+
+  pull_request:
+    types: [opened, synchronize, reopened]
+    paths: [
+      '.github/workflows/fusion.yml',
+      'ggml/**',
+      'tests/fusion/**',
+      'tests/test-fusion.cpp',
+      'tests/test-llama-archs.cpp',
+      'src/models/**'
+    ]
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }}
+  cancel-in-progress: true
+
+env:
+  GGML_NLOOP: 3
+  GGML_N_THREADS: 1
+  LLAMA_ARG_LOG_COLORS: 1
+  LLAMA_ARG_LOG_PREFIX: 1
+  LLAMA_ARG_LOG_TIMESTAMPS: 1
+
+jobs:
+  # TODO: add jobs for other backends as they adopt the fusion debug API
+  metal:
+    runs-on: [self-hosted, macOS, ARM64]
+
+    steps:
+      - name: Clone
+        id: checkout
+        uses: actions/checkout@v6
+
+      - name: Build
+        id: cmake_build
+        run: |
+          cmake -B build \
+            -DCMAKE_BUILD_TYPE=Release \
+            -DLLAMA_FATAL_WARNINGS=ON \
+            -DLLAMA_OPENSSL=OFF \
+            -DGGML_SCHED_NO_REALLOC=ON \
+            -DGGML_BLAS=OFF \
+            -DGGML_METAL=ON
+          time cmake --build build --config Release --target test-llama-archs -j $(sysctl -n hw.logicalcpu)
+          time cmake --build build --config Release --target test-fusion -j $(sysctl -n hw.logicalcpu)
+
+      - name: Generate models
+        id: generate_models
+        run: |
+          rm -rf build-ci-models && mkdir -p build-ci-models
+          ./build/bin/test-llama-archs -o build-ci-models
+
+      - name: Test fusion
+        id: test_fusion
+        run: |
+          ./build/bin/test-fusion --models build-ci-models --device MTL0 --check tests/fusion/MTL.csv
@@ -1,445 +0,0 @@
-name: Models Backend Check
-
-on:
-  workflow_dispatch: # allows manual triggering
-  push:
-    branches:
-      - master
-    paths: [
-      '.github/workflows/models-check.yml',
-      'ggml/**',
-      'tests/fusion/**',
-      'tests/test-fusion.cpp',
-      'tests/test-llama-archs.cpp',
-      'src/llama-graph.cpp',
-      'src/llama-model*',
-      'src/models/**'
-    ]
-
-  pull_request:
-    types: [opened, synchronize, reopened]
-    paths: [
-      '.github/workflows/models-check.yml',
-      'ggml/**',
-      'tests/fusion/**',
-      'tests/test-fusion.cpp',
-      'tests/test-llama-archs.cpp',
-      'src/llama-graph.cpp',
-      'src/llama-model*',
-      'src/models/**'
-    ]
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }}
-  cancel-in-progress: true
-
-env:
-  # note: this is dud token to avoid rate limiting (https://github.com/ggml-org/llama.cpp/pull/25706#issuecomment-4979941302)
-  HF_TOKEN: ${{ secrets.HF_TOKEN_CI }}
-  GGML_NLOOP: 3
-  GGML_N_THREADS: 1
-  LLAMA_ARG_LOG_COLORS: 1
-  LLAMA_ARG_LOG_PREFIX: 1
-  LLAMA_ARG_LOG_TIMESTAMPS: 1
-
-jobs:
-  cuda:
-    runs-on: "hf-jobs-t4-medium:cuda13"
-
-    steps:
-      - name: Clone
-        id: checkout
-        uses: actions/checkout@v6
-
-      - name: Install dependencies
-        run: |
-          sudo apt update
-          sudo apt install -y cmake time python3 python3-venv python3-pip
-
-      - name: ccache
-        uses: ggml-org/ccache-action@v1.2.24
-        with:
-          restore: false
-          save: false
-
-      - name: ccache-buckets-restore
-        uses: ./.github/actions/ccache-buckets
-        with:
-          key: models-check-cuda
-          folder: llama.cpp
-          hf_bucket: ggml-org/cache
-
-      - name: Build
-        id: cmake_build
-        run: |
-          cmake -B build \
-            -DCMAKE_BUILD_TYPE=Release \
-            -DLLAMA_FATAL_WARNINGS=ON \
-            -DLLAMA_OPENSSL=OFF \
-            -DGGML_SCHED_NO_REALLOC=ON \
-            -DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc \
-            -DGGML_CUDA=ON
-          time cmake --build build --config Release --target test-llama-archs -j$(nproc)
-          time cmake --build build --config Release --target test-fusion -j$(nproc)
-
-      - name: ccache-buckets-save
-        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
-        uses: ./.github/actions/ccache-buckets
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
-        with:
-          key: models-check-cuda
-          folder: llama.cpp
-          evict-old-files: 1d
-          hf_bucket: ggml-org/cache
-          save: true
-
-      # - name: Generate models
-      #   id: generate_models
-      #   run: |
-      #     rm -rf build-ci-models && mkdir -p build-ci-models
-      #     ./build/bin/test-llama-archs -o build-ci-models
-
-      # TODO: add for backends as they adopt the fusion debug API
-      # - name: Test fusion
-      #   id: test_fusion
-      #   run: |
-      #     ./build/bin/test-fusion --models build-ci-models --device CUDA0 --check tests/fusion/CUDA.csv
-
-      - name: Test archs
-        id: test_archs
-        run: |
-          GGML_CUDA_DEVICES=1 ./build/bin/test-llama-archs -s 1
-          GGML_CUDA_DEVICES=2 ./build/bin/test-llama-archs -s 1
-          GGML_CUDA_DEVICES=3 ./build/bin/test-llama-archs -s 1
-          GGML_CUDA_DEVICES=4 ./build/bin/test-llama-archs -s 1
-
-  metal:
-    runs-on: [self-hosted, macOS, ARM64]
-
-    steps:
-      - name: Clone
-        id: checkout
-        uses: actions/checkout@v6
-
-      - name: Build
-        id: cmake_build
-        run: |
-          cmake -B build \
-            -DCMAKE_BUILD_TYPE=Release \
-            -DLLAMA_FATAL_WARNINGS=ON \
-            -DLLAMA_OPENSSL=OFF \
-            -DGGML_SCHED_NO_REALLOC=ON \
-            -DGGML_BLAS=OFF \
-            -DGGML_METAL=ON
-          time cmake --build build --config Release --target test-llama-archs -j $(sysctl -n hw.logicalcpu)
-          time cmake --build build --config Release --target test-fusion -j $(sysctl -n hw.logicalcpu)
-
-      - name: Generate models
-        id: generate_models
-        run: |
-          rm -rf build-ci-models && mkdir -p build-ci-models
-          ./build/bin/test-llama-archs -o build-ci-models
-
-      - name: Test fusion
-        id: test_fusion
-        run: |
-          ./build/bin/test-fusion --models build-ci-models --device MTL0 --check tests/fusion/MTL.csv
-
-      - name: Test archs
-        id: test_archs
-        run: |
-          GGML_METAL_DEVICES=1 ./build/bin/test-llama-archs -s 1
-          GGML_METAL_DEVICES=2 ./build/bin/test-llama-archs -s 1
-          GGML_METAL_DEVICES=3 ./build/bin/test-llama-archs -s 1
-          GGML_METAL_DEVICES=4 ./build/bin/test-llama-archs -s 1
-
-  rocm:
-    runs-on: [self-hosted, Linux, gfx1201]
-    container: "rocm/dev-ubuntu-24.04:7.2.4-complete"
-
-    steps:
-      - name: Clone
-        id: checkout
-        uses: actions/checkout@v6
-
-      - name: Install dependencies
-        run: |
-          apt update
-          apt install -y build-essential jq cmake time python3 python3-venv python3-pip
-
-      - name: ccache
-        uses: ggml-org/ccache-action@v1.2.24
-        with:
-          restore: false
-          save: false
-
-      - name: ccache-buckets-restore
-        uses: ./.github/actions/ccache-buckets
-        with:
-          key: models-check-rocm
-          folder: llama.cpp
-          hf_bucket: ggml-org/cache
-
-      - name: Build
-        id: cmake_build
-        run: |
-          cmake -B build \
-            -DCMAKE_BUILD_TYPE=Release \
-            -DLLAMA_FATAL_WARNINGS=ON \
-            -DLLAMA_OPENSSL=OFF \
-            -DGGML_SCHED_NO_REALLOC=ON \
-            -DCMAKE_HIP_COMPILER=$(hipconfig -l)/clang \
-            -DGPU_TARGETS=gfx1201 \
-            -DGGML_HIP=ON
-          time cmake --build build --config Release --target test-llama-archs -j$(nproc)
-          time cmake --build build --config Release --target test-fusion -j$(nproc)
-
-      - name: ccache-buckets-save
-        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
-        uses: ./.github/actions/ccache-buckets
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
-        with:
-          key: models-check-rocm
-          folder: llama.cpp
-          evict-old-files: 1d
-          hf_bucket: ggml-org/cache
-          save: true
-
-      # - name: Generate models
-      #   id: generate_models
-      #   run: |
-      #     rm -rf build-ci-models && mkdir -p build-ci-models
-      #     ./build/bin/test-llama-archs -o build-ci-models
-
-      # TODO: add for backends as they adopt the fusion debug API
-      # - name: Test fusion
-      #   id: test_fusion
-      #   run: |
-      #     ./build/bin/test-fusion --models build-ci-models --device CUDA0 --check tests/fusion/CUDA.csv
-
-      - name: Test archs
-        id: test_archs
-        run: |
-          GGML_CUDA_DEVICES=1 ./build/bin/test-llama-archs -s 1
-          GGML_CUDA_DEVICES=2 ./build/bin/test-llama-archs -s 1
-          GGML_CUDA_DEVICES=3 ./build/bin/test-llama-archs -s 1
-          GGML_CUDA_DEVICES=4 ./build/bin/test-llama-archs -s 1
-
-  vulkan-nvidia:
-    runs-on: "hf-jobs-t4-small:ubuntu26_04"
-
-    steps:
-      - name: Clone
-        id: checkout
-        uses: actions/checkout@v6
-
-      - name: Install dependencies
-        run: |
-          sudo apt update
-          sudo apt install -y build-essential cmake libxcb-xinput0 libxcb-xinerama0 libxcb-cursor-dev libvulkan-dev glslc spirv-headers vulkan-tools mesa-vulkan-drivers libglvnd0 libgl1 libglx0 libegl1 libgles2 time python3 python3-venv python3-pip
-
-      - name: ccache
-        uses: ggml-org/ccache-action@v1.2.24
-        with:
-          restore: false
-          save: false
-
-      - name: ccache-buckets-restore
-        uses: ./.github/actions/ccache-buckets
-        with:
-          key: models-check-vulkan-nvidia
-          folder: llama.cpp
-          hf_bucket: ggml-org/cache
-
-      - name: Build
-        id: cmake_build
-        run: |
-          cmake -B build \
-            -DCMAKE_BUILD_TYPE=Release \
-            -DLLAMA_FATAL_WARNINGS=ON \
-            -DLLAMA_OPENSSL=OFF \
-            -DGGML_SCHED_NO_REALLOC=ON \
-            -DGGML_VULKAN=ON
-          time cmake --build build --config Release --target test-llama-archs -j$(nproc)
-          time cmake --build build --config Release --target test-fusion -j$(nproc)
-
-      - name: ccache-buckets-save
-        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
-        uses: ./.github/actions/ccache-buckets
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
-        with:
-          key: models-check-vulkan-nvidia
-          folder: llama.cpp
-          evict-old-files: 1d
-          hf_bucket: ggml-org/cache
-          save: true
-
-      # - name: Generate models
-      #   id: generate_models
-      #   run: |
-      #     rm -rf build-ci-models && mkdir -p build-ci-models
-      #     ./build/bin/test-llama-archs -o build-ci-models
-
-      # TODO: add for backends as they adopt the fusion debug API
-      # - name: Test fusion
-      #   id: test_fusion
-      #   run: |
-      #     ./build/bin/test-fusion --models build-ci-models --device Vulkan0 --check tests/fusion/Vulkan.csv
-
-      - name: Test archs
-        id: test_archs
-        run: |
-          ./build/bin/test-llama-archs -s 1
-
-  vulkan-amd:
-    runs-on: [self-hosted, Linux, gfx1201]
-    container: "ubuntu:26.04"
-
-    steps:
-      - name: Clone
-        id: checkout
-        uses: actions/checkout@v6
-
-      - name: Install dependencies
-        run: |
-          apt update
-          apt install -y build-essential jq cmake libxcb-xinput0 libxcb-xinerama0 libxcb-cursor-dev libvulkan-dev glslc spirv-headers vulkan-tools mesa-vulkan-drivers libglvnd0 libgl1 libglx0 libegl1 libgles2 time python3 python3-venv python3-pip
-
-      - name: ccache
-        uses: ggml-org/ccache-action@v1.2.24
-        with:
-          restore: false
-          save: false
-
-      - name: ccache-buckets-restore
-        uses: ./.github/actions/ccache-buckets
-        with:
-          key: models-check-vulkan-amd
-          folder: llama.cpp
-          hf_bucket: ggml-org/cache
-
-      - name: Build
-        id: cmake_build
-        run: |
-          cmake -B build \
-            -DCMAKE_BUILD_TYPE=Release \
-            -DLLAMA_FATAL_WARNINGS=ON \
-            -DLLAMA_OPENSSL=OFF \
-            -DGGML_SCHED_NO_REALLOC=ON \
-            -DGGML_VULKAN=ON
-          time cmake --build build --config Release --target test-llama-archs -j$(nproc)
-          time cmake --build build --config Release --target test-fusion -j$(nproc)
-
-      - name: ccache-buckets-save
-        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
-        uses: ./.github/actions/ccache-buckets
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
-        with:
-          key: models-check-vulkan-amd
-          folder: llama.cpp
-          evict-old-files: 1d
-          hf_bucket: ggml-org/cache
-          save: true
-
-      # - name: Generate models
-      #   id: generate_models
-      #   run: |
-      #     rm -rf build-ci-models && mkdir -p build-ci-models
-      #     ./build/bin/test-llama-archs -o build-ci-models
-
-      # TODO: add for backends as they adopt the fusion debug API
-      # - name: Test fusion
-      #   id: test_fusion
-      #   run: |
-      #     ./build/bin/test-fusion --models build-ci-models --device Vulkan0 --check tests/fusion/Vulkan.csv
-
-      - name: Test archs
-        id: test_archs
-        run: |
-          ./build/bin/test-llama-archs -s 1
-
-  webgpu-nvidia:
-    runs-on: "hf-jobs-t4-small:ubuntu26_04"
-
-    steps:
-      - name: Clone
-        id: checkout
-        uses: actions/checkout@v6
-
-      - name: Install dependencies
-        run: |
-          sudo apt update
-          sudo apt install -y build-essential cmake libxcb-xinput0 libxcb-xinerama0 libxcb-cursor-dev libvulkan1 mesa-vulkan-drivers libglvnd0 libgl1 libglx0 libegl1 libgles2 time python3 python3-venv python3-pip
-
-      - name: ccache
-        uses: ggml-org/ccache-action@v1.2.24
-        with:
-          restore: false
-          save: false
-
-      - name: ccache-buckets-restore
-        uses: ./.github/actions/ccache-buckets
-        with:
-          key: models-check-webgpu-nvidia
-          folder: llama.cpp
-          hf_bucket: ggml-org/cache
-
-      - name: Dawn Dependency
-        id: dawn-depends
-        run: |
-          DAWN_VERSION="v20260908.214631"
-          DAWN_OWNER="google"
-          DAWN_REPO="dawn"
-          DAWN_ASSET_NAME="Dawn-94c3c9cc0d5fb2e85aebb370fa8d37b71aa34655-ubuntu-latest-Release"
-          echo "Fetching release asset from https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
-          curl -L -o artifact.tar.gz \
-            "https://github.com/google/dawn/releases/download/${DAWN_VERSION}/${DAWN_ASSET_NAME}.tar.gz"
-          mkdir dawn
-          tar -xvf artifact.tar.gz -C dawn --strip-components=1
-
-      - name: Build
-        id: cmake_build
-        run: |
-          cmake -B build \
-            -DCMAKE_BUILD_TYPE=Release \
-            -DLLAMA_FATAL_WARNINGS=ON \
-            -DLLAMA_OPENSSL=OFF \
-            -DGGML_SCHED_NO_REALLOC=ON \
-            -DCMAKE_PREFIX_PATH="$GITHUB_WORKSPACE/dawn" \
-            -DDawn_DIR="$GITHUB_WORKSPACE/dawn/lib64/cmake/Dawn" \
-            -DGGML_WEBGPU=ON
-          time cmake --build build --config Release --target test-llama-archs -j$(nproc)
-          time cmake --build build --config Release --target test-fusion -j$(nproc)
-
-      - name: ccache-buckets-save
-        if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
-        uses: ./.github/actions/ccache-buckets
-        env:
-          HF_TOKEN: ${{ secrets.HF_TOKEN_CACHE_OUTPUT }}
-        with:
-          key: models-check-webgpu-nvidia
-          folder: llama.cpp
-          evict-old-files: 1d
-          hf_bucket: ggml-org/cache
-          save: true
-
-      # - name: Generate models
-      #   id: generate_models
-      #   run: |
-      #     rm -rf build-ci-models && mkdir -p build-ci-models
-      #     ./build/bin/test-llama-archs -o build-ci-models
-
-      # TODO: add for backends as they adopt the fusion debug API
-      # - name: Test fusion
-      #   id: test_fusion
-      #   run: |
-      #     ./build/bin/test-fusion --models build-ci-models --device WebGPU --check tests/fusion/WebGPU.csv
-
-      - name: Test archs
-        id: test_archs
-        run: |
-          ./build/bin/test-llama-archs -s 1
@@ -323,22 +323,21 @@ jobs:
        include:
          # label = short version used in artifact names / release body
          # cuda  = full container image tag
-          # CTK >= 13.5 bundles CCCL >= 3.5; omit GGML_CUDA_CCCL_VERSION for those versions.
          - build: 'x64'
            os: ubuntu-24.04
            cuda: '12.8.2'
            label: '12.8'
-            defines: '-DGGML_CUDA_CCCL_VERSION=v3.4.3'
+            defines: '-DGGML_CUDA_CUB_3DOT2=ON'
          - build: 'x64'
            os: ubuntu-24.04
            cuda: '13.4.1'
            label: '13.4'
-            defines: '-DGGML_CUDA_CCCL_VERSION=v3.4.3'
+            defines: ''
          - build: 'arm64'
            os: ubuntu-24.04-arm
            cuda: '13.4.1'
            label: '13.4'
-            defines: '-DGGML_CUDA_CCCL_VERSION=v3.4.3'
+            defines: ''

    runs-on: ${{ matrix.os }}
    container: nvidia/cuda:${{ matrix.cuda }}-devel-ubuntu24.04
@@ -1245,16 +1244,15 @@ jobs:
    strategy:
      matrix:
        include:
-          # CTK >= 13.5 bundles CCCL >= 3.5; omit GGML_CUDA_CCCL_VERSION for those versions.
          - cuda: '12.4'
            arch: x64
-            defines: '-DGGML_CUDA_CCCL_VERSION=v3.4.3'
+            defines: '-DGGML_CUDA_CUB_3DOT2=ON'
          - cuda: '13.4'
            arch: x64
-            defines: '-DGGML_CUDA_CCCL_VERSION=v3.4.3'
+            defines: ''
          - cuda: '13.4'
            arch: arm64
-            defines: '-DGGML_CUDA_CCCL_VERSION=v3.4.3 -DCMAKE_TOOLCHAIN_FILE=cmake/arm64-windows-msvc-cuda.cmake'
+            defines: '-DCMAKE_TOOLCHAIN_FILE=cmake/arm64-windows-msvc-cuda.cmake'

    steps:
      - name: Clone
@@ -1281,6 +1279,7 @@ jobs:
      - name: Build
        id: cmake_build
        shell: cmd
+        # TODO: Remove GGML_CUDA_CUB_3DOT2 flag once CCCL 3.2 is bundled within CTK and that CTK version is used in this project
        run: |
          call "C:\Program Files\Microsoft Visual Studio\2022\Enterprise\VC\Auxiliary\Build\vcvarsall.bat" ${{ matrix.arch == 'x64' && 'x64' || 'amd64_arm64' }}
          cmake -S . -B build -G "Ninja Multi-Config" ^
@@ -1345,11 +1344,11 @@ jobs:
        shell: bash

    env:
-      WINDOWS_BASEKIT_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/0cb67a0d-67f6-410b-868b-f4a0a17ff0cf/intel-oneapi-toolkit-2026.1.1.32_offline.exe
+      WINDOWS_BASEKIT_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/b60765d1-2b85-4e85-86b6-cb0e9563a699/intel-deep-learning-essentials-2025.3.3.18_offline.exe
      WINDOWS_DPCPP_MKL: intel.oneapi.win.cpp-dpcpp-common:intel.oneapi.win.mkl.devel:intel.oneapi.win.dnnl:intel.oneapi.win.tbb.devel
-      LEVEL_ZERO_SDK_URL: https://github.com/oneapi-src/level-zero/releases/download/v1.33.1/level-zero-win-sdk-1.33.1.zip
+      LEVEL_ZERO_SDK_URL: https://github.com/oneapi-src/level-zero/releases/download/v1.28.2/level-zero-win-sdk-1.28.2.zip
      ONEAPI_ROOT: "C:/Program Files (x86)/Intel/oneAPI"
-      ONEAPI_INSTALLER_VERSION: "2026.1"
+      ONEAPI_INSTALLER_VERSION: "2025.3.3"

    steps:
      - name: Clone
@@ -1392,11 +1391,9 @@ jobs:
        run: |
          echo "cp oneAPI running time dll files in ${{ env.ONEAPI_ROOT }} to ./build/bin"

-          cp "${{ env.ONEAPI_ROOT }}/mkl/latest/bin/mkl_sycl_blas.6.dll" ./build/bin
-          cp "${{ env.ONEAPI_ROOT }}/mkl/latest/bin/mkl_core.3.dll" ./build/bin
-          cp "${{ env.ONEAPI_ROOT }}/mkl/latest/bin/mkl_def.3.dll" ./build/bin
-          cp "${{ env.ONEAPI_ROOT }}/mkl/latest/bin/mkl_avx2.3.dll" ./build/bin
-          cp "${{ env.ONEAPI_ROOT }}/mkl/latest/bin/mkl_tbb_thread.3.dll" ./build/bin
+          cp "${{ env.ONEAPI_ROOT }}/mkl/latest/bin/mkl_sycl_blas.5.dll" ./build/bin
+          cp "${{ env.ONEAPI_ROOT }}/mkl/latest/bin/mkl_core.2.dll" ./build/bin
+          cp "${{ env.ONEAPI_ROOT }}/mkl/latest/bin/mkl_tbb_thread.2.dll" ./build/bin

          cp "${{ env.ONEAPI_ROOT }}/compiler/latest/bin/ur_adapter_level_zero.dll" ./build/bin
          cp "${{ env.ONEAPI_ROOT }}/compiler/latest/bin/ur_adapter_level_zero_v2.dll" ./build/bin
@@ -1411,11 +1408,13 @@ jobs:
            echo "Level Zero loader DLL not found in oneAPI or SDK; relying on system driver/runtime"
          fi

-          cp "${{ env.ONEAPI_ROOT }}/compiler/latest/bin/sycl9.dll" ./build/bin
+          cp "${{ env.ONEAPI_ROOT }}/compiler/latest/bin/sycl8.dll" ./build/bin
          cp "${{ env.ONEAPI_ROOT }}/compiler/latest/bin/svml_dispmd.dll" ./build/bin
          cp "${{ env.ONEAPI_ROOT }}/compiler/latest/bin/libmmd.dll" ./build/bin
          cp "${{ env.ONEAPI_ROOT }}/compiler/latest/bin/libiomp5md.dll" ./build/bin
          cp "${{ env.ONEAPI_ROOT }}/compiler/latest/bin/sycl-ls.exe" ./build/bin
+          cp "${{ env.ONEAPI_ROOT }}/compiler/latest/bin/libsycl-fallback-bfloat16.spv" ./build/bin
+          cp "${{ env.ONEAPI_ROOT }}/compiler/latest/bin/libsycl-native-bfloat16.spv" ./build/bin

          cp "${{ env.ONEAPI_ROOT }}/dnnl/latest/bin/dnnl.dll" ./build/bin
          cp "${{ env.ONEAPI_ROOT }}/tbb/latest/bin/tbb12.dll" ./build/bin
@@ -1455,8 +1454,8 @@ jobs:

    env:
      ONEAPI_ROOT: /opt/intel/oneapi/
-      ONEAPI_INSTALLER_VERSION: "2026.1"
-      LEVEL_ZERO_VERSION: "1.33.1"
+      ONEAPI_INSTALLER_VERSION: "2025.3.3"
+      LEVEL_ZERO_VERSION: "1.28.2"
      LEVEL_ZERO_UBUNTU_VERSION: "u24.04"

    steps:
@@ -1470,16 +1469,16 @@ jobs:
        shell: bash
        run: |
          cd /tmp
-          wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/5996e26b-f48a-42b1-8db0-b002ad0bd8d7/intel-oneapi-toolkit-2026.1.1.33_offline.sh -O intel-oneapi-toolkit_offline.sh
-          sudo bash intel-oneapi-toolkit_offline.sh -s -a --silent --eula accept
+          wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/56f7923a-adb8-43f3-8b02-2b60fcac8cab/intel-deep-learning-essentials-2025.3.3.16_offline.sh -O intel-deep-learning-essentials_offline.sh
+          sudo bash intel-deep-learning-essentials_offline.sh -s -a --silent --eula accept

      - name: Install Level Zero SDK
        shell: bash
        run: |
          cd /tmp
-          wget -q "https://github.com/oneapi-src/level-zero/releases/download/v${LEVEL_ZERO_VERSION}/libze1_${LEVEL_ZERO_VERSION}%2B${LEVEL_ZERO_UBUNTU_VERSION}_amd64.deb" -O libze1.deb
-          wget -q "https://github.com/oneapi-src/level-zero/releases/download/v${LEVEL_ZERO_VERSION}/libze-dev_${LEVEL_ZERO_VERSION}%2B${LEVEL_ZERO_UBUNTU_VERSION}_amd64.deb" -O libze-dev.deb
-          sudo apt-get install -y ./libze1.deb ./libze-dev.deb
+          wget -q "https://github.com/oneapi-src/level-zero/releases/download/v${LEVEL_ZERO_VERSION}/level-zero_${LEVEL_ZERO_VERSION}%2B${LEVEL_ZERO_UBUNTU_VERSION}_amd64.deb" -O level-zero.deb
+          wget -q "https://github.com/oneapi-src/level-zero/releases/download/v${LEVEL_ZERO_VERSION}/level-zero-devel_${LEVEL_ZERO_VERSION}%2B${LEVEL_ZERO_UBUNTU_VERSION}_amd64.deb" -O level-zero-devel.deb
+          sudo apt-get install -y ./level-zero.deb ./level-zero-devel.deb

      - name: Download UI build
        uses: actions/download-artifact@v7
@@ -77,7 +77,7 @@
 /ggml/src/ggml-vulkan/                  @ggml-org/ggml-vulkan
 /ggml/src/ggml-webgpu/                  @ggml-org/ggml-webgpu
 /ggml/src/ggml-zdnn/                    @ggml-org/ggml-zdnn @Andreas-Krebbel @AlekseiNikiforovIBM
-/ggml/src/ggml-zendnn/                  @avinashcpandey @Jiten1parmar
+/ggml/src/ggml-zendnn/                  @avinashcpandey @Jiten1parmar @z-vishal
 /ggml/src/ggml.c                        @ggerganov
 /ggml/src/ggml.cpp                      @ggerganov
 /ggml/src/gguf.cpp                      @JohannesGaessler @Green-Sky
@@ -21,13 +21,13 @@ docker run --privileged -it \
    -v $HOME/llama.cpp/ci-cache:/ci-cache \
    -v $HOME/llama.cpp/ci-results:/ci-results \
    -v $PWD:/ws -w /ws \
-    registry.mthreads.com/mcconline/musa_sdk:5.2.0-devel-ubuntu22.04-s5000
+    registry.mthreads.com/mcconline/inference/pytorch:2.9.1.post1-py3.10-musa5.2.0-mp31-devel-ubuntu22.04-amd64
 ```

 Inside the container, execute the following commands:

 ```bash
-apt update -y && apt install -y bc cmake ccache git python3.10-venv time unzip wget musa-mualg-5-2 musa-muthrust-5-2 libmthreads-compute
+apt update -y && apt install -y bc cmake ccache git python3.10-venv time unzip wget
 git config --global --add safe.directory /ws
 GG_BUILD_MUSA=1 bash ./ci/run.sh /ci-results /ci-cache
 ```
@@ -49,6 +49,14 @@ mkdir -p "$2"
 OUT=$(realpath "$1")
 MNT=$(realpath "$2")

+# gpu-rocm self-hosted runner can't upload logs to blob; keep each run's logs in
+# their own dir keyed by the GitHub run id so an Actions run URL maps to its logs.
+if [ -n "${GG_BUILD_ROCM}" ] && [ -n "${GITHUB_RUN_ID}" ]; then
+    OUT="$OUT/run-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT:-1}"
+    mkdir -p "$OUT"
+    echo "ci results dir: $OUT"
+fi
+
 rm -f $OUT/*.log

 sd=`dirname $0`
@@ -72,8 +80,8 @@ else
 fi

 if [ ! -z ${GG_BUILD_CUDA} ]; then
-    # TODO: Drop GGML_CUDA_CCCL_VERSION when CUDA CI uses CTK >= 13.5, which bundles CCCL >= 3.5.
-    CMAKE_EXTRA="${CMAKE_EXTRA} -DGGML_CUDA=ON -DGGML_CUDA_CCCL_VERSION=v3.4.3"
+    # TODO: Remove GGML_CUDA_CUB_3DOT2 flag once CCCL 3.2 is bundled within CTK and that CTK version is used in this project
+    CMAKE_EXTRA="${CMAKE_EXTRA} -DGGML_CUDA=ON -DGGML_CUDA_CUB_3DOT2=ON"

    if command -v nvidia-smi >/dev/null 2>&1; then
        CUDA_ARCH=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader,nounits 2>/dev/null | head -1 | tr -d '.')
@@ -351,7 +351,7 @@ static bool parse_bool_value(const std::string & value) {
 static std::string get_default_local_path(const std::string & url) {
    auto f = string_split<std::string>(url, '#').front();
    f = string_split<std::string>(f, '?').front();
-    return fs_path_to_utf8(fs_get_cache_file(string_split<std::string>(f, '/').back()));
+    return fs_get_cache_file(string_split<std::string>(f, '/').back());
 }

 static bool spec_types_is_default(const common_params & params) {
@@ -387,9 +387,6 @@ common_models_handler common_models_handler_init(const common_params & params, l
            break;
        }
    }
-    if (curr_ex == LLAMA_EXAMPLE_DOWNLOAD) {
-        use_mmproj = true;
-    }

    opts.bearer_token    = params.hf_token;
    opts.offline         = params.offline;
@@ -720,24 +717,24 @@ void common_models_handler_apply(common_models_handler & handler, common_params
 // 1. system-wide: /etc/llama.cpp/config.ini (%PROGRAMDATA%\llama.cpp\config.ini on windows)
 // 2. user-level: ${XDG_CONFIG_HOME:-~/.config}/llama.cpp/config.ini (%APPDATA%\llama.cpp\config.ini on windows)
 static void common_params_apply_system_config(common_params & params, llama_example ex) {
-    std::vector<std::filesystem::path> paths;
+    std::vector<std::string> paths;

 #if defined(_WIN32)
-    const std::filesystem::path program_data = common_get_path_from_env("PROGRAMDATA");
+    const std::string program_data = common_get_env("PROGRAMDATA");
    if (!program_data.empty()) {
-        paths.push_back(program_data / "llama.cpp" / "config.ini");
+        paths.push_back(program_data + "\\llama.cpp\\config.ini");
    }
 #else
    paths.push_back("/etc/llama.cpp/config.ini");
 #endif

    try {
-        paths.push_back(fs_get_config_directory() / "config.ini");
+        paths.push_back(fs_get_config_directory() + "config.ini");
    } catch (const std::exception & e) {
        LOG_DBG("cannot read user-level config file, skipping: %s\n", e.what());
    }

-    std::vector<std::filesystem::path> found;
+    std::vector<std::string> found;
    for (const auto & path : paths) {
        std::error_code ec;
        if (std::filesystem::exists(path, ec)) {
@@ -751,7 +748,7 @@ static void common_params_apply_system_config(common_params & params, llama_exam
    common_preset_context ctx(ex);
    ctx.ignore_unknown_keys = true; // the same config file is shared by all programs
    for (const auto & path : found) {
-        LOG_INF("using config file: %s\n", fs_path_to_utf8(path).c_str());
+        LOG_INF("using config file: %s\n", path.c_str());
        common_preset global;
        common_presets presets = ctx.load_from_ini(path, global);
        global.apply_to_params(params);
@@ -1099,12 +1099,6 @@ std::optional<common_chat_params> common_chat_try_specialized_template(
        return common_chat_params_init_ministral_3(tmpl, params);
    }

-    // LLM-jp-4.1 - GPT-OSS dialect (spaces after special tokens, <|end|>-separated parallel calls)
-    if (src.find("chat_format=llm-jp-harmony-v1") != std::string::npos) {
-        LOG_DBG("Using specialized template: LLM-jp Harmony v1\n");
-        return common_chat_params_init_llm_jp_harmony(tmpl, params);
-    }
-
    // GPT-OSS - has unique channel-based structure that needs dedicated handler
    if (src.find("<|channel|>") != std::string::npos) {
        LOG_DBG("Using specialized template: GPT-OSS\n");
@@ -49,7 +49,7 @@
 #include <unistd.h>
 #endif

-#if !defined(_WIN32)
+#if defined(__linux__)
 #include <sys/types.h>
 #include <pwd.h>
 #endif
@@ -613,6 +613,34 @@ std::string string_from(const struct llama_context * ctx, const std::vector<llam
    return buf.str();
 }

+std::string string_from(const struct llama_context * ctx, const struct llama_batch & batch) {
+    std::stringstream buf;
+
+    buf << "[ ";
+
+    bool first = true;
+    for (int i = 0; i < batch.n_tokens; ++i) {
+        if (!first) {
+            buf << ", ";
+        } else {
+            first = false;
+        }
+
+        auto detokenized = common_token_to_piece(ctx, batch.token[i]);
+
+        buf << "\n"          << std::to_string(i)
+            << ", token '"   << detokenized << "'"
+            << ", pos "      << std::to_string(batch.pos[i])
+            << ", n_seq_id " << std::to_string(batch.n_seq_id[i])
+            << ", seq_id "   << std::to_string(batch.seq_id[i][0])
+            << ", logits "   << std::to_string(batch.logits[i]);
+    }
+
+    buf << " ]";
+
+    return buf.str();
+}
+
 void string_process_escapes(std::string & input) {
    std::size_t input_len = input.length();
    std::size_t output_idx = 0;
@@ -912,27 +940,11 @@ std::string fs_path_to_utf8(const std::filesystem::path & path) {
    return std::string(value.begin(), value.end());
 }

-void fs_write_atomic(const std::filesystem::path & path, const std::string & data) {
+// returns true if successful, false otherwise
+bool fs_create_directory_with_parents(const std::string & path) {
    std::error_code ec;
-    std::filesystem::path path_tmp = path;
-    path_tmp += ".tmp";
-
-    if (path.has_parent_path()) {
-        std::filesystem::create_directories(path.parent_path(), ec);
-    }
-
-    std::ofstream file(path_tmp, std::ios::binary);
-    file << data;
-    file.close();
-
-    if (!file.fail()) {
-        std::filesystem::rename(path_tmp, path, ec);
-    }
-
-    if (file.fail() || ec) {
-        std::filesystem::remove(path_tmp, ec);
-        throw std::runtime_error("failed to write file: " + fs_path_to_utf8(path));
-    }
+    std::filesystem::create_directories(std::filesystem::u8path(path), ec);
+    return !ec;
 }

 bool fs_is_directory(const std::string & path) {
@@ -957,77 +969,113 @@ void common_set_env(const std::string & name, const std::string & value) {
 #endif
 }

-std::filesystem::path common_get_path_from_env(const std::string & name) {
-#if defined(_WIN32)
-    const std::wstring wname = utf8_to_wstring(name);
-    const wchar_t * wvalue = _wgetenv(wname.c_str());
-    return wvalue ? std::filesystem::path(wvalue) : std::filesystem::path();
-#else
-    const char * value = std::getenv(name.c_str());
-    return value ? std::filesystem::path(value) : std::filesystem::path();
-#endif
-}
-
-#if !defined(_WIN32)
-static std::filesystem::path get_home_directory() {
-    std::filesystem::path home = common_get_path_from_env("HOME");
-    if (!home.empty()) {
-        return home;
-    }
-    const struct passwd * pw = getpwuid(getuid());
-    if (!pw || !pw->pw_dir || !*pw->pw_dir) {
-        throw std::runtime_error("Failed to find $HOME directory");
-    }
-    return pw->pw_dir;
-}
-#endif
-
-std::filesystem::path fs_get_cache_directory() {
-    std::filesystem::path cache_directory = common_get_path_from_env("LLAMA_CACHE");
-    if (!cache_directory.empty()) {
-        return cache_directory;
-    }
-#if defined(_WIN32)
-    cache_directory = common_get_path_from_env("LOCALAPPDATA");
+std::string fs_get_cache_directory() {
+    std::string cache_directory = "";
+    auto ensure_trailing_slash = [](std::string p) {
+        // Make sure to add trailing slash
+        if (p.empty() || p.back() != DIRECTORY_SEPARATOR) {
+            p += DIRECTORY_SEPARATOR;
+        }
+        return p;
+    };
+    cache_directory = common_get_env("LLAMA_CACHE");
    if (cache_directory.empty()) {
-        throw std::runtime_error("Failed to find %LOCALAPPDATA% directory");
-    }
+#if defined(__linux__) || defined(__FreeBSD__) || defined(_AIX) || \
+        defined(__OpenBSD__) || defined(__NetBSD__)
+        const std::string xdg_cache_home = common_get_env("XDG_CACHE_HOME");
+        const std::string home           = common_get_env("HOME");
+        if (!xdg_cache_home.empty()) {
+            cache_directory = xdg_cache_home;
+        } else if (!home.empty()) {
+            cache_directory = home + "/.cache/";
+        } else {
+#if defined(__linux__)
+            /* no $HOME is defined, fallback to getpwuid */
+            struct passwd *pw = getpwuid(getuid());
+            if ((!pw) || (!pw->pw_dir)) {
+                throw std::runtime_error("Failed to find $HOME directory");
+            }
+
+            cache_directory = std::string(pw->pw_dir) + std::string("/.cache/");
+#else /* defined(__linux__) */
+            throw std::runtime_error("Failed to find $HOME directory");
+#endif /* defined(__linux__) */
+        }
 #elif defined(__APPLE__)
-    cache_directory = get_home_directory() / "Library/Caches";
+        cache_directory = common_get_env("HOME");
+        if (cache_directory.empty()) {
+            throw std::runtime_error("Failed to find $HOME directory");
+        }
+        cache_directory += "/Library/Caches/";
+#elif defined(_WIN32)
+        cache_directory = common_get_env("LOCALAPPDATA");
+        if (cache_directory.empty()) {
+            throw std::runtime_error("Failed to find %LOCALAPPDATA% directory");
+        }
+#elif defined(__EMSCRIPTEN__)
+        GGML_ABORT("not implemented on this platform");
 #else
-    cache_directory = common_get_path_from_env("XDG_CACHE_HOME");
-    if (cache_directory.empty()) {
-        cache_directory = get_home_directory() / ".cache";
-    }
+#  error Unknown architecture
 #endif
-    return cache_directory / "llama.cpp";
+        cache_directory = ensure_trailing_slash(cache_directory);
+        cache_directory += "llama.cpp";
+    }
+    return ensure_trailing_slash(cache_directory);
 }

-std::filesystem::path fs_get_config_directory() {
-    std::filesystem::path config_directory;
-#if defined(_WIN32)
-    config_directory = common_get_path_from_env("APPDATA");
+std::string fs_get_config_directory() {
+    std::string config_directory = "";
+    auto ensure_trailing_slash = [](std::string p) {
+        if (p.empty() || p.back() != DIRECTORY_SEPARATOR) {
+            p += DIRECTORY_SEPARATOR;
+        }
+        return p;
+    };
+#if defined(__linux__) || defined(__FreeBSD__) || defined(_AIX) || \
+        defined(__OpenBSD__) || defined(__NetBSD__) || defined(__APPLE__)
+    const std::string xdg_config_home = common_get_env("XDG_CONFIG_HOME");
+    const std::string home            = common_get_env("HOME");
+    if (!xdg_config_home.empty()) {
+        config_directory = xdg_config_home;
+    } else if (!home.empty()) {
+        config_directory = home + "/.config/";
+    } else {
+#if defined(__linux__)
+        /* no $HOME is defined, fallback to getpwuid */
+        struct passwd *pw = getpwuid(getuid());
+        if ((!pw) || (!pw->pw_dir)) {
+            throw std::runtime_error("Failed to find $HOME directory");
+        }
+
+        config_directory = std::string(pw->pw_dir) + std::string("/.config/");
+#else
+        throw std::runtime_error("Failed to find $HOME directory");
+#endif
+    }
+#elif defined(_WIN32)
+    config_directory = common_get_env("APPDATA");
    if (config_directory.empty()) {
        throw std::runtime_error("Failed to find %APPDATA% directory");
    }
+#elif defined(__EMSCRIPTEN__)
+    // caller decides what to do when there is no config directory
+    throw std::runtime_error("not implemented on this platform");
 #else
-    config_directory = common_get_path_from_env("XDG_CONFIG_HOME");
-    if (config_directory.empty()) {
-        config_directory = get_home_directory() / ".config";
-    }
+#  error Unknown architecture
 #endif
-    return config_directory / "llama.cpp";
+    config_directory = ensure_trailing_slash(config_directory);
+    config_directory += "llama.cpp";
+    return ensure_trailing_slash(config_directory);
 }

-std::filesystem::path fs_get_cache_file(const std::string & filename) {
+std::string fs_get_cache_file(const std::string & filename) {
    GGML_ASSERT(filename.find(DIRECTORY_SEPARATOR) == std::string::npos);
-    const std::filesystem::path cache_directory = fs_get_cache_directory();
-    std::error_code ec;
-    std::filesystem::create_directories(cache_directory, ec);
-    if (ec) {
-        throw std::runtime_error("failed to create cache directory: " + fs_path_to_utf8(cache_directory));
+    std::string cache_directory = fs_get_cache_directory();
+    const bool success = fs_create_directory_with_parents(cache_directory);
+    if (!success) {
+        throw std::runtime_error("failed to create cache directory: " + cache_directory);
    }
-    return cache_directory / std::filesystem::u8path(filename);
+    return cache_directory + filename;
 }

 std::vector<common_file_info> fs_list(const std::string & path, bool include_directories) {
@@ -1432,8 +1480,7 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode
        }

        if (llama_model_has_encoder(model)) {
-            common_batch batch = common_batch_get_one(lctx, tmp);
-            llama_process(lctx, LLAMA_PROCESS_TYPE_ENCODE, batch.get());
+            llama_encode(lctx, llama_batch_get_one(tmp.data(), tmp.size()));
            llama_token decoder_start_token_id = llama_model_decoder_start_token(model);
            if (decoder_start_token_id == LLAMA_TOKEN_NULL) {
                decoder_start_token_id = bos;
@@ -1442,9 +1489,7 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode
            tmp.push_back(decoder_start_token_id);
        }
        if (llama_model_has_decoder(model)) {
-            tmp.resize(std::min(tmp.size(), (size_t) params.n_batch));
-            common_batch batch = common_batch_get_one(lctx, tmp);
-            llama_process(lctx, LLAMA_PROCESS_TYPE_DECODE, batch.get());
+            llama_decode(lctx, llama_batch_get_one(tmp.data(), std::min(tmp.size(), (size_t) params.n_batch)));
        }
        llama_memory_clear(llama_get_memory(lctx), true);
        llama_synchronize(lctx);
@@ -1508,13 +1553,9 @@ common_context_seq_rm_type common_context_can_seq_rm(llama_context * ctx) {
    tmp.push_back(0);
    tmp.push_back(0);

-    int ret;
-    {
-        common_batch batch = common_batch_get_one(ctx, tmp);
-        ret = llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get());
-    }
+    int ret = llama_decode(ctx, llama_batch_get_one(tmp.data(), tmp.size()));
    if (ret != 0) {
-        COM_ERR("llama_process() failed: %d\n", ret);
+        COM_ERR("llama_decode() failed: %d\n", ret);
        res = COMMON_CONTEXT_SEQ_RM_TYPE_NO;
        goto done;
    }
@@ -1738,6 +1779,33 @@ void common_threadpools::init(llama_context * ctx, const common_params & params)
    llama_attach_threadpool(ctx, threadpool, threadpool_batch);
 }

+//
+// Batch utils
+//
+
+void common_batch_clear(struct llama_batch & batch) {
+    batch.n_tokens = 0;
+}
+
+void common_batch_add(
+                 struct llama_batch & batch,
+                        llama_token   id,
+                          llama_pos   pos,
+    const std::vector<llama_seq_id> & seq_ids,
+                               bool   logits) {
+    GGML_ASSERT(batch.seq_id[batch.n_tokens] && "llama_batch size exceeded");
+
+    batch.token   [batch.n_tokens] = id;
+    batch.pos     [batch.n_tokens] = pos;
+    batch.n_seq_id[batch.n_tokens] = seq_ids.size();
+    for (size_t i = 0; i < seq_ids.size(); ++i) {
+        batch.seq_id[batch.n_tokens][i] = seq_ids[i];
+    }
+    batch.logits  [batch.n_tokens] = logits;
+
+    batch.n_tokens++;
+}
+
 //
 // Vocab utils
 //
@@ -2074,132 +2142,32 @@ float lr_opt::get_lr(float epoch) const {
 }

 bool common_replay_last_token(struct llama_context * ctx, llama_token last_token, int32_t pos) {
-    common_batch batch(ctx);
-    batch.add(last_token, pos, 0, true);
-
-    if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
+    llama_batch batch = llama_batch_get_one(&last_token, 1);
+    batch.pos = &pos;
+    if (llama_decode(ctx, batch)) {
        LOG_ERR("%s: failed to replay last token\n", __func__);
        return false;
    }
    return true;
 }

-common_batch::common_batch(llama_context * ctx) : batch(llama_batch_ext_init(ctx)) {
-    const auto rope_type = llama_model_rope_type(llama_get_model(ctx));
-    n_pos = rope_type == LLAMA_ROPE_TYPE_MROPE || rope_type == LLAMA_ROPE_TYPE_IMROPE ? GGML_MROPE_SECTIONS : 1;
-}
-
-void common_batch::clear() {
-    tokens.clear();
-}
-
-int32_t common_batch::add(llama_token id, llama_pos pos, llama_seq_id seq_id, bool output) {
-    tokens.push_back({ id, { pos, 0, 0, 0 }, seq_id, output, { nullptr, 0, 0 }, {} });
-    return size() - 1;
-}
-
-int32_t common_batch::add(llama_token id, llama_pos pos, const std::vector<llama_seq_id> & seq_ids, bool output) {
-    GGML_ASSERT(!seq_ids.empty());
-
-    const int32_t idx = add(id, pos, seq_ids[0], output);
-    for (size_t s = 1; s < seq_ids.size(); ++s) {
-        add_seq(idx, seq_ids[s]);
-    }
-    return idx;
-}
-
-bool common_batch::add_seq(int32_t idx, llama_seq_id seq_id) {
-    if (idx < 0 || idx >= size()) {
-        return false;
-    }
-    tokens[idx].seq_ids_extra.push_back(seq_id);
-    return true;
-}
-
-bool common_batch::set_output(int32_t idx, bool value) {
-    if (idx < 0 || idx >= size()) {
-        return false;
-    }
-    tokens[idx].output = value;
-    return true;
-}
-
-bool common_batch::set_embd(int32_t idx, llama_embd embd) {
-    if (idx < 0 || idx >= size() || tokens[idx].embd.data != nullptr) {
-        return false;
-    }
-    tokens[idx].embd = embd;
-    return true;
-}
-
-int32_t common_batch::add_embd(llama_embd embd, const llama_pos * pos, llama_seq_id seq_id, bool output) {
-    token t = { LLAMA_TOKEN_NULL, { 0, 0, 0, 0 }, seq_id, output, embd, {} };
-    for (int32_t j = 0; j < n_pos; ++j) {
-        t.pos[j] = pos[j];
-    }
-    tokens.push_back(t);
-    return size() - 1;
-}
-
-llama_batch_ext * common_batch::get_sub_batch(int32_t off, int32_t n) {
-    GGML_ASSERT(batch && "common_batch was not initialized with a context");
-    GGML_ASSERT(off >= 0 && n >= 0 && off + n <= size());
-
-    llama_batch_ext * res = batch.get();
-    llama_batch_ext_clear(res);
-
-    for (int32_t i = off; i < off + n; ++i) {
-        const token & t = tokens[i];
-
-        int32_t idx;
-        if (t.id != LLAMA_TOKEN_NULL) {
-            idx = llama_batch_ext_add_token(res, t.seq_id, t.id);
-            if (idx < 0) {
-                GGML_ABORT("%s: failed to add token %d at index %d (error %d, n = %d)\n", __func__, t.id, i, idx, n);
-            }
-            llama_batch_ext_set_pos(res, idx, t.pos.data());
-            if (t.embd.data && !llama_batch_ext_set_embd_token(res, idx, t.embd)) {
-                GGML_ABORT("%s: failed to set the embedding of token %d at index %d\n", __func__, t.id, i);
-            }
-        } else {
-            idx = llama_batch_ext_add_embd(res, t.seq_id, t.embd);
-            if (idx < 0) {
-                GGML_ABORT("%s: failed to add embedding at index %d (error %d, n = %d)\n", __func__, i, idx, n);
-            }
-            llama_batch_ext_set_pos(res, idx, t.pos.data());
-        }
-        GGML_ASSERT(idx == i - off);
-
-        for (const llama_seq_id seq_id : t.seq_ids_extra) {
-            if (!llama_batch_ext_add_seq(res, idx, seq_id)) {
-                GGML_ABORT("%s: failed to add seq %d to the entry at index %d\n", __func__, seq_id, i);
-            }
-        }
-        if (t.output) {
-            llama_batch_ext_set_output_logits(res, idx, true);
-        }
-    }
-
-    return res;
-}
-
-common_batch common_batch_get_one(llama_context * ctx, const llama_token * tokens, int32_t n_tokens) {
-    common_batch batch(ctx);
+llama_batch_ext_ptr common_batch_ext_get_one(llama_context * ctx, const llama_tokens & tokens) {
+    llama_batch_ext_ptr batch(llama_batch_ext_init(ctx));

    auto mem = llama_get_memory(ctx);
-    llama_pos pos = llama_memory_seq_pos_max(mem, 0) + 1; // -1 + 1 == 0 when the memory is empty
+    llama_pos pos = mem ? llama_memory_seq_pos_max(mem, 0) + 1 : 0;

-    for (int32_t i = 0; i < n_tokens; ++i) {
-        const bool output = i == n_tokens - 1;
-        batch.add(tokens[i], pos, 0, output);
+    for (size_t i = 0; i < tokens.size(); ++i) {
+        const int32_t idx = llama_batch_ext_add_token(batch.get(), 0, tokens[i]);
+        llama_batch_ext_set_pos(batch.get(), idx, &pos);
        pos++;
    }

-    return batch;
-}
+    if (!tokens.empty()) {
+        llama_batch_ext_set_output_logits(batch.get(), (int32_t) tokens.size() - 1, true);
+    }

-common_batch common_batch_get_one(llama_context * ctx, const llama_tokens & tokens) {
-    return common_batch_get_one(ctx, tokens.data(), (int32_t) tokens.size());
+    return batch;
 }

 bool common_prompt_batch_decode(
@@ -2226,7 +2194,7 @@ bool common_prompt_batch_decode(
        // memory, so we can't just remove the last token from the memory and replay the last token which
        // is the reason for this logic.
        llama_tokens prefix_tokens(all_tokens.begin() + offset, all_tokens.begin() + offset + n_tokens_before_last);
-        common_batch batch_prefix = common_batch_get_one(ctx, prefix_tokens);
+        llama_batch_ext_ptr batch_prefix = common_batch_ext_get_one(ctx, prefix_tokens);
        if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch_prefix.get())) {
            COM_ERR("%s", "failed to eval\n");
            return false;
@@ -2236,8 +2204,10 @@ bool common_prompt_batch_decode(
        llama_state_save_file(ctx, state_path.data(), all_tokens.data(), all_tokens.size());
        COM_INF("saved session before last token to %s, n_new = %zu\n", state_path.data(), all_tokens.size());

-        common_batch batch_last(ctx);
-        batch_last.add(all_tokens.back(), n_past, 0, true);
+        llama_token last_token = all_tokens.back();
+        llama_batch_ext_ptr batch_last = common_batch_ext_get_one(ctx, { last_token });
+        llama_pos pos = n_past;
+        llama_batch_ext_set_pos(batch_last.get(), 0, &pos);

        if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch_last.get())) {
            COM_ERR("%s", "failed to eval last token\n");
@@ -2246,7 +2216,7 @@ bool common_prompt_batch_decode(
        n_past++;
    } else {
        llama_tokens new_tokens(all_tokens.begin() + offset, all_tokens.begin() + offset + n_new);
-        common_batch batch = common_batch_get_one(ctx, new_tokens);
+        llama_batch_ext_ptr batch = common_batch_ext_get_one(ctx, new_tokens);
        if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
            COM_ERR("%s", "failed to eval\n");
            return false;
@@ -8,7 +8,6 @@
 #include "ggml.h"
 #include "llama.h"

-#include <array>
 #include <list>
 #include <set>
 #include <sstream>
@@ -880,6 +879,7 @@ void string_process_escapes(std::string & input);
 std::string string_from(bool value);
 std::string string_from(const std::vector<int> & values);
 std::string string_from(const struct llama_context * ctx, const std::vector<llama_token> & tokens);
+std::string string_from(const struct llama_context * ctx, const struct llama_batch & batch);

 bool glob_match(const std::string & pattern, const std::string & str);

@@ -904,19 +904,17 @@ std::string fs_path_to_utf8(const std::filesystem::path & path);
 std::string common_get_env(const std::string & name);
 void        common_set_env(const std::string & name, const std::string & value);

-// reads a path from the environment, an unset variable gives an empty path
-std::filesystem::path common_get_path_from_env(const std::string & name);
-
 //
 // Filesystem utils
 //

 bool fs_validate_filename(const std::string & filename, bool allow_subdirs = false);
+bool fs_create_directory_with_parents(const std::string & path);
 bool fs_is_directory(const std::string & path);

-std::filesystem::path fs_get_cache_directory();
-std::filesystem::path fs_get_cache_file(const std::string & filename);
-std::filesystem::path fs_get_config_directory();
+std::string fs_get_cache_directory();
+std::string fs_get_cache_file(const std::string & filename);
+std::string fs_get_config_directory();

 struct common_file_info {
    std::string path;
@@ -929,8 +927,6 @@ std::vector<common_file_info> fs_list(const std::string & path, bool include_dir
 // fs open, also handle UTF8 on Windows
 std::ifstream fs_open_ifstream(const std::string & fname, std::ios_base::openmode mode);

-void fs_write_atomic(const std::filesystem::path & path, const std::string & data);
-
 //
 // TTY utils
 //
@@ -1031,62 +1027,18 @@ struct common_memory {
 // Batch utils
 //

-// wrapper around llama_batch_ext that provide getter functions for downstream code
-// entries can exceed n_batch, use get_sub_batch() to decode them in chunks
-struct common_batch {
-    struct token {
-        llama_token  id;
-        std::array<llama_pos, GGML_MROPE_SECTIONS> pos; // only pos[0] is used for text tokens
-        llama_seq_id seq_id; // the first sequence id, see add_seq()
-        bool         output;
-        llama_embd   embd; // non-owning view of the data passed to add_embd()/set_embd(), data == NULL if none
-        std::vector<llama_seq_id> seq_ids_extra; // see add_seq()
-    };
+void common_batch_clear(struct llama_batch & batch);

-    std::vector<token> tokens; // mirror of the entries, tokens[i] describes batch index i
-    llama_batch_ext_ptr batch;
-
-    int32_t n_pos = 1; // positions per embedding entry, GGML_MROPE_SECTIONS for MROPE/IMROPE
-
-    common_batch() = default;
-    common_batch(struct llama_context * ctx);
-
-    llama_batch_ext * get() { return get_sub_batch(0, size()); }
-
-    // render entries [off, off + n) into batch, the result is overwritten by the next call
-    llama_batch_ext * get_sub_batch(int32_t off, int32_t n);
-
-    // content type of the batch, all entries carry the same combination
-    bool has_token() const { return !tokens.empty() && tokens[0].id != LLAMA_TOKEN_NULL; }
-    bool has_embd () const { return !tokens.empty() && tokens[0].embd.data != nullptr; }
-
-    void clear();
-
-    // returns the batch index
-    int32_t add(llama_token id, llama_pos pos, llama_seq_id seq_id, bool output);
-
-    // same, with the entry shared by all seq_ids (must not be empty)
-    int32_t add(llama_token id, llama_pos pos, const std::vector<llama_seq_id> & seq_ids, bool output);
-
-    // add the entry at idx to another sequence, tokens[idx].seq_id keeps the first one
-    bool add_seq(int32_t idx, llama_seq_id seq_id);
-
-    bool set_output(int32_t idx, bool value);
-
-    // attach a token embedding to the entry at idx, can only be set once per entry
-    bool set_embd(int32_t idx, llama_embd embd);
-
-    // add an embedding-only entry (no token id)
-    // pos points to n_pos positions
-    int32_t add_embd(llama_embd embd, const llama_pos * pos, llama_seq_id seq_id, bool output);
-
-    int32_t size() const { return (int32_t) tokens.size(); }
-};
+void common_batch_add(
+                 struct llama_batch & batch,
+                        llama_token   id,
+                          llama_pos   pos,
+    const std::vector<llama_seq_id> & seq_ids,
+                               bool   logits);

 // create a single-sequence batch from a list of tokens
-// positions continue from the memory, last token always have output_logits set to true
-common_batch common_batch_get_one(struct llama_context * ctx, const llama_token * tokens, int32_t n_tokens);
-common_batch common_batch_get_one(struct llama_context * ctx, const llama_tokens & tokens);
+// last token always have output_logits set to true
+llama_batch_ext_ptr common_batch_ext_get_one(struct llama_context * ctx, const llama_tokens & tokens);

 // decodes a single batch of tokens for a prompt and manages session tokens
 //
@@ -1019,7 +1019,6 @@ namespace console {
                    line.clear();
                    pop_cursor();
                }
-                line += '\n';
                has_more = false;
            }
        } else {
@@ -1051,6 +1050,7 @@ namespace console {
        if (!std::getline(std::wcin, wline)) {
            // Input stream is bad or EOF received
            line.clear();
+            GenerateConsoleCtrlEvent(CTRL_C_EVENT, 0);
            return false;
        }

@@ -1065,7 +1065,7 @@ namespace console {
        if (!line.empty()) {
            char last = line.back();
            if (last == '/') { // Always return control on '/' symbol
-                line.back() = '\n';
+                line.pop_back();
                return false;
            }
            if (last == '\\') { // '\\' changes the default action
@@ -46,9 +46,39 @@
 // downloader
 //

+// validate repo name format: owner/repo
+static void write_file(const std::string & fname, const std::string & content) {
+    const std::string fname_tmp = fname + ".tmp";
+    std::ofstream     file(fname_tmp);
+    if (!file) {
+        throw std::runtime_error(string_format("error: failed to open file '%s'\n", fname.c_str()));
+    }
+
+    try {
+        file << content;
+        file.close();
+
+        // Makes write atomic
+        if (rename(fname_tmp.c_str(), fname.c_str()) != 0) {
+            LOG_ERR("%s: unable to rename file: %s to %s\n", __func__, fname_tmp.c_str(), fname.c_str());
+            // If rename fails, try to delete the temporary file
+            if (remove(fname_tmp.c_str()) != 0) {
+                LOG_ERR("%s: unable to delete temporary file: %s\n", __func__, fname_tmp.c_str());
+            }
+        }
+    } catch (...) {
+        // If anything fails, try to delete the temporary file
+        if (remove(fname_tmp.c_str()) != 0) {
+            LOG_ERR("%s: unable to delete temporary file: %s\n", __func__, fname_tmp.c_str());
+        }
+
+        throw std::runtime_error(string_format("error: failed to write file '%s'\n", fname.c_str()));
+    }
+}
+
 static void write_etag(const std::string & path, const std::string & etag) {
    const std::string etag_path = path + ".etag";
-    fs_write_atomic(std::filesystem::u8path(etag_path), etag);
+    write_file(etag_path, etag);
    LOG_DBG("%s: file etag saved: %s\n", __func__, etag_path.c_str());
 }

@@ -244,12 +274,6 @@ static bool common_pull_file(httplib::Client & cli,
        return false;
    }

-    ofs.close();
-    if (!ofs) {
-        LOG_ERR("%s: error closing file: %s\n", __func__, path_tmp.c_str());
-        return false;
-    }
-
    return true;
 }

@@ -262,7 +286,7 @@ static int common_download_file_single_online(const std::string & url,
    static const int max_attempts        = 3;
    static const int retry_delay_seconds = 2;

-    const bool file_exists = std::filesystem::exists(std::filesystem::u8path(path));
+    const bool file_exists = std::filesystem::exists(path);

    if (file_exists && skip_etag) {
        LOG_DBG("%s: using cached file: %s\n", __func__, path.c_str());
@@ -453,7 +477,7 @@ int common_download_file_single(const std::string & url,
        return common_download_file_single_online(url, path, online_opts, skip_etag);
    }

-    if (!std::filesystem::exists(std::filesystem::u8path(path))) {
+    if (!std::filesystem::exists(path)) {
        LOG_ERR("%s: required file is not available in cache (offline mode): %s\n", __func__, path.c_str());
        return -1;
    }
@@ -919,7 +943,7 @@ std::string common_docker_resolve_model(const std::string & docker) {
        std::string model_filename = repo;
        std::replace(model_filename.begin(), model_filename.end(), '/', '_');
        model_filename += "_" + tag + ".gguf";
-        std::string local_path = fs_path_to_utf8(fs_get_cache_file(model_filename));
+        std::string local_path = fs_get_cache_file(model_filename);

        const std::string blob_url = url_prefix + "/blobs/" + gguf_digest;
        common_download_opts opts;
@@ -44,7 +44,8 @@ static fs::path get_cache_directory() {
            {HOME_DIR,                fs::path(".cache") / "huggingface" / "hub"}
        };
        for (const auto & entry : entries) {
-            if (fs::path base = common_get_path_from_env(entry.var); !base.empty()) {
+            if (auto * p = std::getenv(entry.var); p && *p) {
+                fs::path base(p);
                return entry.path.empty() ? base : base / entry.path;
            }
        }
@@ -172,6 +173,28 @@ static bool is_valid_subpath(const fs::path & path, const fs::path & subpath) {
    return b_end == b.end();
 }

+static void safe_write_file(const fs::path & path, const std::string & data) {
+    fs::path path_tmp = path.string() + ".tmp";
+
+    if (path.has_parent_path()) {
+        fs::create_directories(path.parent_path());
+    }
+
+    std::ofstream file(path_tmp);
+    file << data;
+    file.close();
+
+    std::error_code ec;
+
+    if (!file.fail()) {
+        fs::rename(path_tmp, path, ec);
+    }
+    if (file.fail() || ec) {
+        fs::remove(path_tmp, ec);
+        throw std::runtime_error("failed to write file: " + path.string());
+    }
+}
+
 static common_json api_get(const std::string & url,
                           const std::string & token) {
    auto [cli, parts] = common_http_client(url);
@@ -218,7 +241,6 @@ static std::string get_repo_commit(const std::string & repo_id,
        fs::path refs_path = get_repo_path(repo_id) / "refs";
        std::string name;
        std::string commit;
-        fs::path name_path;

        for (const auto & branch : json["branches"]) {
            if (!branch.is_object() ||
@@ -229,28 +251,24 @@ static std::string get_repo_commit(const std::string & repo_id,
            std::string _name = branch["name"].get<std::string>();
            std::string _commit = branch["targetCommit"].get<std::string>();

-            if (!is_valid_commit(_commit)) {
-                LOG_WRN("%s: skip invalid commit: %s\n", __func__, _commit.c_str());
+            if (!is_valid_subpath(refs_path, _name)) {
+                LOG_WRN("%s: skip invalid branch: %s\n", __func__, _name.c_str());
                continue;
            }
-            const fs::path candidate = fs::u8path(_name);
-
-            if (!is_valid_subpath(refs_path, candidate)) {
-                LOG_WRN("%s: skip invalid branch: %s\n", __func__, _name.c_str());
+            if (!is_valid_commit(_commit)) {
+                LOG_WRN("%s: skip invalid commit: %s\n", __func__, _commit.c_str());
                continue;
            }

            if (_name == "main") {
                name = _name;
                commit = _commit;
-                name_path = candidate;
                break;
            }

            if (name.empty() || commit.empty()) {
                name = _name;
                commit = _commit;
-                name_path = candidate;
            }
        }

@@ -259,7 +277,7 @@ static std::string get_repo_commit(const std::string & repo_id,
            return {};
        }

-        fs_write_atomic(refs_path / name_path, commit);
+        safe_write_file(refs_path / name, commit);
        return commit;

    } catch (const common_json_error & e) {
@@ -308,9 +326,7 @@ hf_files get_repo_files(const std::string & repo_id,
            file.repo_id = repo_id;
            file.path = item["path"].get<std::string>();

-            const fs::path subpath = fs::u8path(file.path);
-
-            if (!is_valid_subpath(commit_path, subpath)) {
+            if (!is_valid_subpath(commit_path, file.path)) {
                LOG_WRN("%s: skip invalid path: %s\n", __func__, file.path.c_str());
                continue;
            }
@@ -330,12 +346,12 @@ hf_files get_repo_files(const std::string & repo_id,

            file.url = endpoint + repo_id + "/resolve/" + commit + "/" + file.path;

-            fs::path final_path = commit_path / subpath;
-            file.final_path = fs_path_to_utf8(final_path);
+            fs::path final_path = commit_path / file.path;
+            file.final_path = final_path.string();

            if (!file.oid.empty() && !fs::exists(final_path)) {
                fs::path local_path = blobs_path / file.oid;
-                file.local_path = fs_path_to_utf8(local_path);
+                file.local_path = local_path.string();
            } else {
                file.local_path = file.final_path;
            }
@@ -402,7 +418,7 @@ hf_files get_cached_files(const std::string & repo_id) {
        if (!fs::exists(snapshots_path)) {
            continue;
        }
-        std::string _repo_id = folder_name_to_repo(fs_path_to_utf8(repo.path().filename()));
+        std::string _repo_id = folder_name_to_repo(repo.path().filename().string());

        if (!is_valid_repo_id(_repo_id)) {
            continue;
@@ -425,9 +441,8 @@ hf_files get_cached_files(const std::string & repo_id) {
            if (!path.empty()) {
                hf_file file;
                file.repo_id = _repo_id;
-                const auto generic_path = path.generic_u8string();
-                file.path = std::string(generic_path.begin(), generic_path.end());
-                file.local_path = fs_path_to_utf8(entry.path());
+                file.path = path.generic_string();
+                file.local_path = entry.path().string();
                file.final_path = file.local_path;
                files.push_back(std::move(file));
            }
@@ -441,8 +456,8 @@ std::string finalize_file(const hf_file & file) {
    static std::atomic<bool> symlinks_disabled{false};

    std::error_code ec;
-    fs::path local_path = fs::u8path(file.local_path);
-    fs::path final_path = fs::u8path(file.final_path);
+    fs::path local_path(file.local_path);
+    fs::path final_path(file.final_path);

    if (local_path == final_path || fs::exists(final_path, ec)) {
        return file.final_path;
@@ -489,7 +504,7 @@ bool remove_cached_repo(const std::string & repo_id) {
    std::error_code ec;
    auto removed = fs::remove_all(repo_path, ec);
    if (ec) {
-        LOG_ERR("%s: failed to remove repo cache %s: %s\n", __func__, fs_path_to_utf8(repo_path).c_str(), ec.message().c_str());
+        LOG_ERR("%s: failed to remove repo cache %s: %s\n", __func__, repo_path.string().c_str(), ec.message().c_str());
        return false;
    }
    return removed > 0;
@@ -537,6 +537,8 @@ value for_statement::execute_impl(context & ctx) const {

    std::vector<value> filtered_items;
    for (size_t i = 0; i < items.size(); ++i) {
+        context loop_scope(scope);
+
        value current = items[i];

        std::function<void(context&)> scope_update_fn = [](context &) { /* no-op */};
@@ -582,7 +584,6 @@ value for_statement::execute_impl(context & ctx) const {
        }

        if (select_expr && test_expr) {
-            context loop_scope(scope);
            scope_update_fn(loop_scope);
            value test_val = test_expr->execute(loop_scope);
            if (!test_val->as_bool()) {
@@ -887,7 +888,7 @@ value member_expression::execute_impl(context & ctx) const {
        JJ_DEBUG("Accessed property '%s' value, got type: %s", key.c_str(), val->type().c_str());

    } else if (is_val<value_array>(object) || is_val<value_string>(object)) {
-        if (is_val<value_int>(property) || is_val<value_bool>(property)) {
+        if (is_val<value_int>(property)) {
            int64_t index = property->as_int();
            JJ_DEBUG("Accessing %s index %d", object->type().c_str(), (int)index);
            if (is_val<value_array>(object)) {
@@ -910,6 +911,8 @@ value member_expression::execute_impl(context & ctx) const {
            JJ_DEBUG("Accessing %s built-in '%s'", is_val<value_array>(object) ? "array" : "string", key.c_str());
            val = try_builtin_func(ctx, key, object, true);

+        } else {
+            throw std::runtime_error("Cannot access property with non-string/non-number: got " + property->type());
        }
    } else {
        if (!is_val<value_string>(property)) {
@@ -923,10 +926,10 @@ value member_expression::execute_impl(context & ctx) const {
        value_t::stats_t::mark_used(val);
        value_t::stats_t::mark_used(object);
        value_t::stats_t::mark_used(property);
-        if (is_val<value_object>(object) || is_val<value_string>(property) || is_val<value_float>(property) || is_val<value_array>(property) || is_val<value_none>(property)) {
-            object->stats.ops.insert("object_access");
-        } else if (is_val<value_int>(property) || is_val<value_bool>(property)) {
+        if (is_val<value_int>(property)) {
            object->stats.ops.insert("array_access");
+        } else if (is_val<value_string>(property)) {
+            object->stats.ops.insert("object_access");
        }
    }

@@ -149,13 +149,6 @@ static value test_type_fn(const func_args & args) {
    JJ_DEBUG("test_type_fn: type=%s, %s or %s result=%d", typeid(T).name(), typeid(U).name(), typeid(V).name(), is_type ? 1 : 0);
    return mk_val<value_bool>(is_type);
 }
-template<typename T, typename U, typename V, typename W>
-static value test_type_fn(const func_args & args) {
-    args.ensure_count(1);
-    bool is_type = is_val<T>(args.get_pos(0)) || is_val<U>(args.get_pos(0)) || is_val<V>(args.get_pos(0)) || is_val<W>(args.get_pos(0));
-    JJ_DEBUG("test_type_fn: type=%s, %s, %s or %s result=%d", typeid(T).name(), typeid(U).name(), typeid(V).name(), typeid(W).name(), is_type ? 1 : 0);
-    return mk_val<value_bool>(is_type);
-}
 template<value_compare_op op>
 static value test_compare_fn(const func_args & args) {
    args.ensure_count(2, 2);
@@ -268,30 +261,6 @@ static value tojson(const func_args & args) {
    return mk_val<value_string>(json_str);
 }

-static value & get_attribute(const value & val, const value & attr, value & default_val) {
-    if (!attr->is_undefined()) {
-        if (is_val<value_array>(val)) {
-            value idx = attr;
-
-            if (is_val<value_string>(attr)) {
-                const std::string s = attr->as_string().str();
-                if (!s.empty() && std::all_of(s.begin(), s.end(), [](unsigned char c) { return std::isdigit(c); })) {
-                    try {
-                        idx = mk_val<value_int>(std::stoll(s));
-                    } catch (...) {
-                        idx = mk_val<value_undefined>();
-                    }
-                }
-            }
-
-            return val->at(idx, default_val);
-        } else if (is_val<value_object>(val)) {
-            return val->at(attr, default_val);
-        }
-    }
-    return default_val;
-}
-
 template<bool is_reject>
 static value selectattr(const func_args & args) {
    args.ensure_count(2, 4);
@@ -305,7 +274,10 @@ static value selectattr(const func_args & args) {
    if (args.count() == 2) {
        // example: array | selectattr("active")
        for (const auto & item : arr) {
-            value attr_val = get_attribute(item, attribute, val_default);
+            if (!is_val<value_object>(item)) {
+                throw raised_exception("selectattr: item is not an object");
+            }
+            value attr_val = item->at(attribute, val_default);
            bool is_selected = attr_val->as_bool();
            if constexpr (is_reject) is_selected = !is_selected;
            if (is_selected) out->push_back(item);
@@ -346,7 +318,10 @@ static value selectattr(const func_args & args) {
        }
        auto test_fn = it->second;
        for (const auto & item : arr) {
-            value attr_val = get_attribute(item, attribute, val_default);
+            if (!is_val<value_object>(item)) {
+                throw raised_exception("selectattr: item is not an object");
+            }
+            value attr_val = item->at(attribute, val_default);
            func_args test_args(args.ctx);
            test_args.push_back(attr_val); // attribute value
            test_args.push_back(extra_arg); // extra argument
@@ -503,8 +478,8 @@ const func_builtins & global_builtins() {
        {"test_is_integer", test_type_fn<value_int>},
        {"test_is_float", test_type_fn<value_float>},
        {"test_is_number", test_type_fn<value_int, value_float>},
-        {"test_is_iterable", test_type_fn<value_object, value_array, value_string, value_undefined>},
-        {"test_is_sequence", test_type_fn<value_object, value_array, value_string, value_undefined>},
+        {"test_is_iterable", test_type_fn<value_array, value_string, value_undefined>},
+        {"test_is_sequence", test_type_fn<value_array, value_string, value_undefined>},
        {"test_is_mapping", test_type_fn<value_object>},
        {"test_is_lower", [](const func_args & args) -> value {
            args.ensure_vals<value_string>();
@@ -1093,14 +1068,22 @@ const func_builtins & value_array_t::get_builtins() const {
            }
            value val_delim = args.get_kwarg_or_pos("d",         1);
            value attribute = args.get_kwarg_or_pos("attribute", 2);
-            value undef = mk_val<value_undefined>();
            const auto & arr = args.get_pos(0)->as_array();
+            const bool attr_is_int = is_val<value_int>(attribute);
+            if (!attribute->is_undefined() && !is_val<value_string>(attribute) && !attr_is_int) {
+                throw raised_exception("join() attribute must be string or integer");
+            }
+            const int64_t attr_int = attr_is_int ? attribute->as_int() : 0;
            const std::string delim = val_delim->is_undefined() ? "" : val_delim->as_string().str();
            std::string result;
            for (size_t i = 0; i < arr.size(); ++i) {
                value val_arr = arr[i];
                if (!attribute->is_undefined()) {
-                    val_arr = get_attribute(val_arr, attribute, undef);
+                    if (attr_is_int && is_val<value_array>(val_arr)) {
+                        val_arr = val_arr->at(attr_int);
+                    } else if (!attr_is_int && is_val<value_object>(val_arr)) {
+                        val_arr = val_arr->at(attribute);
+                    }
                }
                if (!is_val<value_string>(val_arr) && !is_val<value_int>(val_arr) && !is_val<value_float>(val_arr)) {
                    throw raised_exception("join() can only join arrays of strings or numerics");
@@ -1132,11 +1115,21 @@ const func_builtins & value_array_t::get_builtins() const {
            }
            value val       = args.get_pos(0);
            value attribute = args.get_kwarg_or_pos("attribute", 1);
+            const bool attr_is_int = is_val<value_int>(attribute);
+            if (!is_val<value_string>(attribute) && !attr_is_int) {
+                throw raised_exception("map: attribute must be string or integer");
+            }
+            const int64_t attr_int = attr_is_int ? attribute->as_int() : 0;
            value default_val = args.get_kwarg("default", mk_val<value_undefined>());
            auto out = mk_val<value_array>();
            auto arr = val->as_array();
            for (const auto & item : arr) {
-                value attr_val = get_attribute(item, attribute, default_val);
+                value attr_val;
+                if (attr_is_int) {
+                    attr_val = is_val<value_array>(item) ? item->at(attr_int, default_val) : default_val;
+                } else {
+                    attr_val = is_val<value_object>(item) ? item->at(attribute, default_val) : default_val;
+                }
                out->push_back(attr_val);
            }
            return is_val<value_tuple>(val) ? mk_val<value_tuple>(std::move(out->as_array())) : out;
@@ -1173,14 +1166,22 @@ const func_builtins & value_array_t::get_builtins() const {
            // FIXME: sorting is currently always case sensitive
            //const bool case_sensitive = val_case->as_bool(); // undefined == false
            const bool reverse = val_reverse->as_bool(); // undefined == false
-            value undef = mk_val<value_undefined>();
+            const bool attr_is_int = is_val<value_int>(attribute);
+            const int64_t attr_int = attr_is_int ? attribute->as_int() : 0;
            std::vector<value> arr = val->as_array(); // copy
            std::sort(arr.begin(), arr.end(),[&](const value & a, const value & b) {
                value val_a = a;
                value val_b = b;
                if (!attribute->is_undefined()) {
-                    val_a = get_attribute(a, attribute, undef);
-                    val_b = get_attribute(b, attribute, undef);
+                    if (attr_is_int && is_val<value_array>(a) && is_val<value_array>(b)) {
+                        val_a = a->at(attr_int);
+                        val_b = b->at(attr_int);
+                    } else if (!attr_is_int && is_val<value_object>(a) && is_val<value_object>(b)) {
+                        val_a = a->at(attribute);
+                        val_b = b->at(attribute);
+                    } else {
+                        throw raised_exception("sort: unsupported object attribute comparison between " + a->type() + " and " + b->type());
+                    }
                }
                return value_compare(val_a, val_b, reverse ? value_compare_op::gt : value_compare_op::lt);
            });
@@ -1198,23 +1199,19 @@ const func_builtins & value_array_t::get_builtins() const {
            args.ensure_vals<value_array>();
            value val_case    = args.get_kwarg_or_pos("case_sensitive", 1);
            value attribute   = args.get_kwarg_or_pos("attribute",      2);
+            if (!attribute->is_undefined()) {
+                throw not_implemented_exception("min: attribute not implemented");
+            }
            // FIXME: min is currently always case sensitive
            (void) val_case;
-            value undef = mk_val<value_undefined>();
            const auto & arr = args.get_pos(0)->as_array();
            if (arr.empty()) {
-                return undef;
+                return mk_val<value_undefined>();
            }
            value result = arr[0];
-            for (const auto & item : arr) {
-                value val_arr = item;
-                value val_cmp = result;
-                if (!attribute->is_undefined()) {
-                    val_arr = get_attribute(val_arr, attribute, undef);
-                    val_cmp = get_attribute(val_cmp, attribute, undef);
-                }
-                if (value_compare(val_arr, val_cmp, value_compare_op::lt)) {
-                    result = item;
+            for (size_t i = 1; i < arr.size(); ++i) {
+                if (value_compare(arr[i], result, value_compare_op::lt)) {
+                    result = arr[i];
                }
            }
            return result;
@@ -1224,23 +1221,19 @@ const func_builtins & value_array_t::get_builtins() const {
            args.ensure_vals<value_array>();
            value val_case    = args.get_kwarg_or_pos("case_sensitive", 1);
            value attribute   = args.get_kwarg_or_pos("attribute",      2);
+            if (!attribute->is_undefined()) {
+                throw not_implemented_exception("max: attribute not implemented");
+            }
            // FIXME: max is currently always case sensitive
            (void) val_case;
-            value undef = mk_val<value_undefined>();
            const auto & arr = args.get_pos(0)->as_array();
            if (arr.empty()) {
-                return undef;
+                return mk_val<value_undefined>();
            }
            value result = arr[0];
-            for (const auto & item : arr) {
-                value val_arr = item;
-                value val_cmp = result;
-                if (!attribute->is_undefined()) {
-                    val_arr = get_attribute(val_arr, attribute, undef);
-                    val_cmp = get_attribute(val_cmp, attribute, undef);
-                }
-                if (value_compare(val_arr, val_cmp, value_compare_op::gt)) {
-                    result = item;
+            for (size_t i = 1; i < arr.size(); ++i) {
+                if (value_compare(arr[i], result, value_compare_op::gt)) {
+                    result = arr[i];
                }
            }
            return result;
@@ -433,12 +433,6 @@ struct value_array_t : public value_t {
        }
        return val_arr[index];
    }
-    virtual value & at(const value & index, value & default_val) override {
-        if (!is_val<value_int>(index) && !is_val<value_bool>(index)) {
-            return default_val;
-        }
-        return at(index->as_int(), default_val);
-    }
    virtual const func_builtins & get_builtins() const override;
    virtual bool is_hashable() const override {
        if (std::all_of(val_arr.begin(), val_arr.end(), [&](auto & val) -> bool {
@@ -1,164 +0,0 @@
-#include "parsers.h"
-
-// LLM-jp-4.1: the GPT-OSS (Harmony) format with two differences
-//  - the tokenizer emits a space after every special token: "<|channel|> analysis<|message|> ..."
-//  - parallel tool calls are consecutive assistant messages, all but the last closed by <|end|>
-common_chat_params common_chat_params_init_llm_jp_harmony(const common_chat_template &          tmpl,
-                                                          const autoparser::generation_params & inputs) {
-    common_chat_params data;
-
-    // Copy reasoning to the "thinking" field as expected by the template
-    auto adjusted_messages = json::array();
-    for (auto msg : inputs.messages) {
-        if (msg.contains("reasoning_content") && msg.at("reasoning_content").is_string()) {
-            msg["thinking"] = msg.at("reasoning_content");
-            if (msg.contains("tool_calls") && msg.at("tool_calls").is_array() && !msg.at("tool_calls").empty()) {
-                msg.erase("content");
-            }
-        }
-        adjusted_messages.push_back(msg);
-    }
-
-    auto prompt = common_chat_template_direct_apply_impl(tmpl, inputs, /* messages_override= */ adjusted_messages);
-
-    // Check if we need to replace the return token with end token during
-    // inference and without generation prompt. For more details see:
-    // https://github.com/ggml-org/llama.cpp/issues/15417
-    if (inputs.is_inference && !inputs.add_generation_prompt) {
-        static constexpr std::string_view return_token = "<|return|>";
-        static constexpr std::string_view end_token    = "<|end|>";
-        if (size_t pos = prompt.rfind(return_token); pos != std::string::npos) {
-            prompt.replace(pos, return_token.length(), end_token);
-        }
-    }
-
-    data.prompt            = prompt;
-    data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs, /* messages_override= */ adjusted_messages);
-    data.message_delimiters = {
-        { COMMON_CHAT_ROLE_ASSISTANT, "<|start|>assistant" },
-        { COMMON_CHAT_ROLE_USER,      "<|start|>user"      },
-        { COMMON_CHAT_ROLE_SYSTEM,    "<|start|>developer" },
-        { COMMON_CHAT_ROLE_SYSTEM,    "<|start|>system"    },
-        { COMMON_CHAT_ROLE_TOOL,      "<|start|>functions" },
-    };
-
-    data.format            = COMMON_CHAT_FORMAT_PEG_NATIVE;
-    data.supports_thinking = true;
-
-    data.thinking_start_tag = "<|channel|>analysis<|message|>";
-    data.thinking_end_tags  = {"<|end|>"};
-
-    // These special tokens are required to parse properly, so we include them
-    // even if parse_tool_calls is false.
-    data.preserved_tokens = {
-        "<|channel|>", "<|constrain|>", "<|message|>", "<|start|>", "<|end|>",
-    };
-
-    // Adjust prompt for continuation
-    if (inputs.has_continuation()) {
-        const auto & msg = inputs.continue_msg;
-
-        data.generation_prompt = "<|start|>assistant<|channel|>analysis<|message|>" + msg.reasoning_content;
-        if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) {
-            data.generation_prompt += "<|end|><|start|>assistant<|channel|>final<|message|>" + msg.render_content();
-        }
-
-        data.prompt += data.generation_prompt;
-    }
-
-    auto has_tools           = inputs.tools.is_array() && !inputs.tools.empty();
-    auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object();
-    auto include_grammar     = has_response_format || (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE);
-    auto extract_reasoning   = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE;
-
-    auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
-        // tokenizer space after special tokens; not p.space() since GBNF `space` allows one space only
-        auto sp          = p.chars("[ ]", 0, -1);
-        auto channel_tag = p.literal("<|channel|>") + sp;
-        // one space only: keep an intentional leading space in the body
-        auto message     = p.literal("<|message|>") + p.optional(p.literal(" "));
-
-        auto start          = p.rule("start", p.literal("<|start|>") + sp + p.literal("assistant"));
-        auto end            = p.rule("end", p.literal("<|end|>"));
-        auto content        = p.rule("message-content", p.until("<|end|>"));
-        auto channel        = channel_tag + (p.literal("commentary") | p.literal("analysis"));
-        auto constrain_type = p.chars("[A-Za-z0-9_-]", 1, -1);
-        auto constraint     = p.optional(p.space() + p.optional(p.literal("<|constrain|>") + sp) + constrain_type);
-
-        auto start_analysis = channel_tag + p.literal("analysis") + message;
-        if (extract_reasoning) {
-            p.rule("analysis", start_analysis + p.reasoning(content) + end);
-        } else {
-            p.rule("analysis", p.content(start_analysis + content + end));
-        }
-
-        auto analysis  = p.ref("analysis");
-        auto preamble  = p.rule("preamble", channel_tag + p.literal("commentary") + message + p.content(content) + end);
-        auto final_msg = p.rule("final", channel_tag + p.literal("final") + message + p.content(content));
-
-        auto any = p.rule("any", preamble | analysis);
-
-        if (has_response_format) {
-            auto response_format = p.rule("response-format",
-                channel_tag + p.literal("final") + constraint + message +
-                p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema)));
-
-            return p.zero_or_more(start + analysis) + start + response_format;
-        }
-
-        if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) {
-            auto tool_choice = p.choice();
-
-            foreach_function(inputs.tools, [&](const json & tool) {
-                const auto & function = tool.at("function");
-                std::string  name     = function.at("name");
-                const auto   params   = common_chat_tool_parameters(function);
-
-                auto func_name = p.literal(" to=functions.") + p.tool_name(p.literal(name));
-                auto args      = p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", params));
-
-                // recipient in role header
-                //   <|start|>assistant to=functions.NAME<|channel|>(commentary|analysis)[constraint]<|message|>ARGS
-                auto tool_in_role = p.tool(p.tool_open(func_name + channel + constraint + message) + args);
-
-                // recipient in channel header
-                //   <|channel|>(commentary|analysis) to=functions.NAME[constraint]<|message|>ARGS
-                auto tool_in_channel = p.tool(p.tool_open(channel + func_name + constraint + message) + args);
-
-                tool_choice |= p.rule("tool-" + name, tool_in_role | tool_in_channel);
-            });
-
-            // parallel calls are separated by <|end|>; inside the trigger rule so the lazy grammar covers all of them
-            auto tool_calls = inputs.parallel_tool_calls
-                ? tool_choice + p.zero_or_more(end + start + tool_choice)
-                : tool_choice;
-            auto tool_call  = p.trigger_rule("tool-call", tool_calls);
-
-            if (inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED) {
-                return p.zero_or_more(start + any) + start + tool_call;
-            }
-
-            return p.zero_or_more(start + any) + start + (tool_call | final_msg);
-        }
-
-        return p.zero_or_more(start + any) + start + final_msg;
-    });
-
-    data.parser = parser.save();
-
-    if (include_grammar) {
-        data.grammar_lazy = !(has_response_format || (has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED));
-        data.grammar      = build_grammar([&](const common_grammar_builder & builder) {
-            parser.build_grammar(builder, data.grammar_lazy);
-        });
-
-        data.grammar_triggers = {
-            { COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN, "^\\s+to$" },
-            { COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN, "^<\\|channel\\|>\\s*(?:commentary|analysis)\\s+to=functions$" },
-            { COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN, "<\\|start\\|>\\s*assistant(\\s+to)" },
-            { COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN, "<\\|start\\|>\\s*assistant(<\\|channel\\|>\\s*(?:commentary|analysis)\\s+to)" }
-        };
-    }
-
-    return data;
-}
@@ -43,10 +43,9 @@ common_chat_params common_chat_params_init_muse_glimmer(const common_chat_templa

    auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE;

-    auto has_tools           = inputs.tools.is_array() && !inputs.tools.empty();
-    auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object();
-    // Constrained grammar whenever tools are offered or a response format is requested.
-    auto include_grammar     = has_response_format || (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE);
+    auto has_tools = inputs.tools.is_array() && !inputs.tools.empty();
+    // Constrained grammar whenever tools are offered.
+    auto include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE;

    auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
        auto start = p.rule("start", p.literal("<|start|>assistant"));
@@ -66,15 +65,6 @@ common_chat_params common_chat_params_init_muse_glimmer(const common_chat_templa
        auto final_msg  = p.rule("final", recipient + p.literal("<|message|>") +
                                              p.content(p.until_one_of({ "<|eot|>", "<|eom|>" })));

-        if (has_response_format) {
-            auto response_json   = p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema));
-            auto response_format = p.rule("response-format",
-                recipient + p.literal("<|message|>") +
-                ((p.literal("```json") + p.space() + response_json + p.space() + p.literal("```")) | response_json));
-
-            return p.zero_or_more(start + analysis) + start + response_format;
-        }
-
        if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) {
            auto string_value = p.ac(
                p.tool_arg_string_value(p.until("</atem:parameter>")) + p.tool_arg_close(p.literal("</atem:parameter>")),
@@ -134,7 +124,7 @@ common_chat_params common_chat_params_init_muse_glimmer(const common_chat_templa
    data.parser = parser.save();

    if (include_grammar) {
-        data.grammar_lazy = !(has_response_format || (has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED));
+        data.grammar_lazy = inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_REQUIRED;
        data.grammar      = build_grammar([&](const common_grammar_builder & builder) {
            parser.build_grammar(builder, data.grammar_lazy);
        });
@@ -68,8 +68,6 @@ common_chat_params common_chat_params_init_ling3(const common_chat_template & tm
 // tool_list_tokens preserves the LFM2 system tool-list markers; LFM2.5 renders without them
 common_chat_params common_chat_params_init_lfm2(const common_chat_template & tmpl, const autoparser::generation_params & inputs, bool tool_list_tokens);

-common_chat_params common_chat_params_init_llm_jp_harmony(const common_chat_template & tmpl, const autoparser::generation_params & inputs);
-
 common_chat_params common_chat_params_init_minicpm5(const common_chat_template & tmpl, const autoparser::generation_params & inputs);

 common_chat_params common_chat_params_init_minimax_m3(const common_chat_template & tmpl, const autoparser::generation_params & inputs);
@@ -13,7 +13,6 @@ set(LLAMA_CHAT_PARSERS_SOURCES
    ${CMAKE_CURRENT_LIST_DIR}/kimi-k3.cpp
    ${CMAKE_CURRENT_LIST_DIR}/ling3.cpp
    ${CMAKE_CURRENT_LIST_DIR}/lfm2.cpp
-    ${CMAKE_CURRENT_LIST_DIR}/llm-jp-harmony.cpp
    ${CMAKE_CURRENT_LIST_DIR}/minicpm5.cpp
    ${CMAKE_CURRENT_LIST_DIR}/minimax-m3.cpp
    ${CMAKE_CURRENT_LIST_DIR}/ministral3.cpp
@@ -167,16 +167,16 @@ void common_preset::apply_to_params(common_params & params, const std::set<std::
    }
 }

-static std::map<std::string, std::map<std::string, std::string>> parse_ini_from_file(const std::filesystem::path & path) {
+static std::map<std::string, std::map<std::string, std::string>> parse_ini_from_file(const std::string & path) {
    std::map<std::string, std::map<std::string, std::string>> parsed;

    if (!std::filesystem::exists(path)) {
-        throw std::runtime_error("preset file does not exist: " + fs_path_to_utf8(path));
+        throw std::runtime_error("preset file does not exist: " + path);
    }

    std::ifstream file(path);
    if (!file.good()) {
-        throw std::runtime_error("failed to open server preset file: " + fs_path_to_utf8(path));
+        throw std::runtime_error("failed to open server preset file: " + path);
    }

    std::string contents((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
@@ -225,7 +225,7 @@ static std::map<std::string, std::map<std::string, std::string>> parse_ini_from_
    common_peg_parse_context ctx(contents);
    const auto result = parser.parse(ctx);
    if (!result.success()) {
-        throw std::runtime_error("failed to parse server config file: " + fs_path_to_utf8(path));
+        throw std::runtime_error("failed to parse server config file: " + path);
    }

    std::string current_section = COMMON_PRESET_DEFAULT_NAME;
@@ -282,7 +282,7 @@ common_preset_context::common_preset_context(llama_example ex)
    key_to_opt = get_map_key_opt(ctx_params);
 }

-common_presets common_preset_context::load_from_ini(const std::filesystem::path & path, common_preset & global) const {
+common_presets common_preset_context::load_from_ini(const std::string & path, common_preset & global) const {
    common_presets out;
    auto ini_data = parse_ini_from_file(path);

@@ -323,7 +323,7 @@ common_presets common_preset_context::load_from_ini(const std::filesystem::path
                }
                LOG_DBG("accepted option: %s = %s\n", key.c_str(), preset.options[opt].c_str());
            } else if (ignore_unknown_keys) {
-                LOG_WRN("ignoring option '%s' from %s: not supported by this program\n", key.c_str(), fs_path_to_utf8(path).c_str());
+                LOG_WRN("ignoring option '%s' from %s: not supported by this program\n", key.c_str(), path.c_str());
            } else {
                throw std::runtime_error(string_format(
                    "option '%s' not recognized in preset '%s'",
@@ -67,7 +67,7 @@ struct common_preset_context {
    common_preset_context(llama_example ex);

    // load presets from INI file
-    common_presets load_from_ini(const std::filesystem::path & path, common_preset & global) const;
+    common_presets load_from_ini(const std::string & path, common_preset & global) const;

    // generate presets from cached models
    common_presets load_from_cache() const;
@@ -681,8 +681,6 @@ std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sample
    std::vector<llama_token> result;
    result.reserve(idxs.size());

-    const llama_vocab * vocab = llama_model_get_vocab(llama_get_model(ctx));
-
    size_t i = 0;
    for (; i < draft.size(); i++) {
        const llama_token id = common_sampler_sample(gsmpl, ctx, idxs[i], grammar_first);
@@ -691,9 +689,7 @@ std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sample

        result.push_back(id);

-        // do not accept draft tokens after an EOG - they are not output but would stay in the context
-        // on replay the last token is from the target and can be EOG, so a trailing EOG is still accepted
-        if (draft[i] != id || (llama_vocab_is_eog(vocab, id) && i + 1 < draft.size())) {
+        if (draft[i] != id) {
            break;
        }
    }
@@ -165,7 +165,7 @@ struct common_speculative_impl {

    virtual void begin(llama_seq_id seq_id, const llama_tokens & prompt) = 0;

-    virtual bool process(const common_batch & batch) = 0;
+    virtual bool process(const llama_batch & batch) = 0;

    virtual void draft(common_speculative_draft_params_vec & dparams) = 0;

@@ -179,11 +179,7 @@ struct common_speculative_impl {
 struct common_speculative_impl_draft_simple : public common_speculative_impl {
    common_params_speculative_draft params;

-    common_batch batch;
-
-    // zero row at the draft input width, stands in for target embeddings the draft cannot read
-    std::vector<float> zeros;
-    bool zeros_warned = false; // the substitution is reported once
+    llama_batch batch;

    std::vector<common_sampler_ptr> smpls;

@@ -198,8 +194,6 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
            throw std::runtime_error("draft-simple requires a draft context");
        }

-        zeros.assign(llama_model_n_embd_inp(llama_get_model(ctx_dft)), 0.0f);
-
        SPC_TRC("%s", "adding speculative implementation 'draft-simple'\n");
        SPC_TRC("- n_max=%d, n_min=%d, p_min=%f\n", this->params.n_max, this->params.n_min, this->params.p_min);
        SPC_TRC("- gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n",
@@ -210,7 +204,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
                ctx_dft ? "yes" : "no",
                common_speculative_get_devices_str(this->params.devices).c_str());

-        batch = common_batch(ctx_dft);
+        batch = llama_batch_init(llama_n_batch(ctx_dft), 0, 1);

        // TODO: optimize or pass from outside?
        // {
@@ -255,46 +249,21 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
        }
    }

+    ~common_speculative_impl_draft_simple() override {
+        llama_batch_free(batch);
+    }
+
    void begin(llama_seq_id /*seq_id*/, const llama_tokens & /*prompt*/) override {
        // noop
    }

-    bool process(const common_batch & batch_in) override {
+    bool process(const llama_batch & batch) override {
        auto * ctx_dft = params.ctx_dft;

-        // copy the entries to a batch owned by the draft context, only the last token is output
-        batch.clear();
-        const int32_t n_tokens = batch_in.size();
-        for (int32_t k = 0; k < n_tokens; ++k) {
-            const auto & t = batch_in.tokens[k];
-            const bool output = k == n_tokens - 1;
-            if (t.id != LLAMA_TOKEN_NULL) {
-                const int32_t idx = batch.add(t.id, t.pos[0], t.seq_id, output);
-                if (t.embd.data) {
-                    batch.set_embd(idx, t.embd);
-                }
-            } else {
-                // mtmd input is projected by the target encoder, a draft with a different width cannot read it
-                // it gets zeros instead, keeping its positions contiguous
-                // ref: https://github.com/ggml-org/llama.cpp/pull/29385#discussion_r4124743243
-                const size_t n_embd = t.embd.n_rows * t.embd.n_embd;
-                const bool same_width = n_embd == zeros.size();
-                if (!same_width && !zeros_warned) {
-                    SPC_WRN("target embeddings of size %zu do not fit the draft input width %zu, "
-                            "the draft receives zero rows for them and drafts after multimodal input will be poor\n",
-                            n_embd, zeros.size());
-                    zeros_warned = true;
-                }
-                const llama_embd embd = same_width ? t.embd : llama_embd{ zeros.data(), 1, zeros.size() };
-                batch.add_embd(embd, t.pos.data(), t.seq_id, output);
-            }
-        }
+        llama_batch batch_dft = batch;
+        batch_dft.logits = nullptr;

-        if (batch.size() == 0) {
-            return true;
-        }
-
-        const int ret = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch.get());
+        const int ret = llama_decode(ctx_dft, batch_dft);

        if (ret != 0) {
            SPC_ERR("failed to decode draft batch, ret = %d\n", ret);
@@ -308,7 +277,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
    void draft(common_speculative_draft_params_vec & dparams) override {
        auto & ctx_dft = params.ctx_dft;

-        batch.clear();
+        common_batch_clear(batch);

        // keep track of which sequences are still drafting
        int n_drafting = 0;
@@ -325,12 +294,12 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
            drafting[seq_id] = true;
            common_sampler_reset(smpls[seq_id].get());

-            batch.add(dp.id_last, dp.pos0, seq_id, true);
+            common_batch_add(batch, dp.id_last, dp.pos0, { seq_id }, true);
        }

-        int ret = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch.get());
+        int ret = llama_decode(ctx_dft, batch);
        if (ret != 0) {
-            SPC_ERR("llama_process returned %d\n", ret);
+            SPC_ERR("llama_decode returned %d\n", ret);
            return;
        }

@@ -339,7 +308,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
        while (n_drafting > 0) {
            int i_batch = 0;

-            batch.clear();
+            common_batch_clear(batch);

            for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
                if (!drafting[seq_id]) {
@@ -384,17 +353,17 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
                    continue;
                }

-                batch.add(id, dp.pos0 + i + 1, seq_id, true);
+                common_batch_add(batch, id, dp.pos0 + i + 1, { seq_id }, true);
            }

-            if (batch.size() == 0) {
+            if (batch.n_tokens == 0) {
                break;
            }

            // evaluate the drafted tokens on the draft model
-            ret = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch.get());
+            ret = llama_decode(ctx_dft, batch);
            if (ret != 0) {
-                SPC_ERR("llama_process[%d] returned %d\n", i, ret);
+                SPC_ERR("llama_decode[%d] returned %d\n", i, ret);
                break;
            }

@@ -454,8 +423,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl {
 //      encoder+decoder on n_accepted+1 rows).
 struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
    common_params_speculative_draft params;
-    common_batch batch;     // decoder input, (token, g_embd) pairs
-    common_batch batch_enc; // encoder input, built from the extracted target features
+    llama_batch batch;

    std::vector<common_sampler_ptr> smpls;

@@ -509,8 +477,11 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
        n_embd_enc = (int32_t) target_layer_ids_n * n_embd_tgt;
        n_layer_tgt = llama_model_n_layer(model_tgt);

-        batch     = common_batch(ctx_dft);
-        batch_enc = common_batch(ctx_dft);
+        const int32_t n_b = (int32_t) llama_n_batch(ctx_dft);
+        batch = llama_batch_init(/*n_tokens=*/ n_b, /*embd=*/ n_embd_dec, /*n_seq_max=*/ 1);
+        // llama_batch_init allocates only one of token/embd; eagle3 decoder needs both.
+        // TODO: fix, how to call without malloc
+        batch.token = (llama_token *) malloc(sizeof(llama_token) * n_b);

        smpls.resize(n_seq);
        for (auto & s : smpls) {
@@ -572,6 +543,12 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
            llama_sampler_free(backend_chains[seq_id]);
        }
        backend_chains.clear();
+
+        if (batch.token != nullptr) {
+            free(batch.token);
+            batch.token = nullptr;
+        }
+        llama_batch_free(batch);
    }

    void begin(llama_seq_id seq_id, const llama_tokens & prompt) override {
@@ -590,16 +567,16 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
        }
    }

-    bool process(const common_batch & batch_in) override {
-        if (batch_in.size() <= 0) {
+    bool process(const llama_batch & batch_in) override {
+        if (batch_in.n_tokens <= 0) {
            return true;
        }

-        if (!batch_in.has_token() || batch_in.has_embd()) {
+        if (batch_in.token == nullptr || batch_in.embd != nullptr) {
            return true;
        }

-        const int32_t n_tokens = batch_in.size();
+        const int32_t n_tokens = batch_in.n_tokens;

        // i_batch_beg[seq] / i_batch_end[seq]: inclusive batch indices of this seq's
        // first/last token in batch_in. Assumes per-seq tokens are contiguous within
@@ -607,7 +584,8 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
        std::vector<int32_t> i_batch_beg(n_seq, -1);
        std::vector<int32_t> i_batch_end(n_seq, -1);
        for (int k = 0; k < n_tokens; ++k) {
-            const llama_seq_id seq_id = batch_in.tokens[k].seq_id;
+            GGML_ASSERT(batch_in.n_seq_id[k] == 1);
+            const llama_seq_id seq_id = batch_in.seq_id[k][0];
            if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) {
                continue;
            }
@@ -641,23 +619,24 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {

        g_embd_buf.resize((size_t) n_tokens * n_embd_dec);

-        // llama_process() requires the full encoder batch to fit in n_ubatch.
+        // llama_encode() requires the full encoder batch to fit in n_ubatch.
        // Allow batch > ubatch: eagle3's per-token encoder can be chunked safely.
        const int32_t n_ubatch_dft = (int32_t) llama_n_ubatch(ctx_dft);
        for (int32_t i = 0; i < n_tokens; i += n_ubatch_dft) {
            const int32_t n_chunk = std::min(n_ubatch_dft, n_tokens - i);

-            // the per-token encoder does not use positions, generate placeholder ones from the memory state
-            batch_enc.clear();
-            llama_pos pos = llama_memory_seq_pos_max(llama_get_memory(ctx_dft), 0) + 1;
-            for (int32_t j = 0; j < n_chunk; ++j) {
-                batch_enc.add_embd({ features_buf.data() + (size_t) (i + j) * n_embd_enc, 1, (size_t) n_embd_enc }, &pos, 0, true);
-                pos++;
-            }
-
-            const int32_t rc = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_ENCODE, batch_enc.get());
+            llama_batch enc_batch = {
+                /*.n_tokens =*/ n_chunk,
+                /*.token    =*/ nullptr,
+                /*.embd     =*/ features_buf.data() + (size_t) i * n_embd_enc,
+                /*.pos      =*/ nullptr,
+                /*.n_seq_id =*/ nullptr,
+                /*.seq_id   =*/ nullptr,
+                /*.logits   =*/ nullptr,
+            };
+            const int32_t rc = llama_encode(ctx_dft, enc_batch);
            if (rc != 0) {
-                SPC_ERR("llama_process(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
+                SPC_ERR("llama_encode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
                        rc, (int) n_chunk, (int) i);
                return false;
            }
@@ -685,7 +664,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
        //       deferred boundary, completed by the next process() or draft() call.
        //   (c) refresh deferred state — stash this ubatch's full g_embd into verify_g,
        //       update pending_g_last / pending_pos_last to the last row.
-        batch.clear();
+        common_batch_clear(batch);

        for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
            const int32_t beg = i_batch_beg[seq_id];
@@ -700,34 +679,36 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
            //   2) pending_pos_last + 1 == pos[beg]
            //   3) pending_pos_last > dft_pos_max // TODO: is this check needed?
            const llama_pos pending_pos = pending_pos_last[seq_id];
-            if (pending_pos >= 0 && pending_pos + 1 == batch_in.tokens[beg].pos[0]) {
+            if (pending_pos >= 0 && pending_pos + 1 == batch_in.pos[beg]) {
                const llama_pos dft_pos_max = llama_memory_seq_pos_max(llama_get_memory(ctx_dft), seq_id);
                if (pending_pos > dft_pos_max) {
-                    const int32_t idx = batch.add(batch_in.tokens[beg].id, pending_pos, seq_id, /*output=*/ false);
-                    batch.set_embd(idx, { pending_g_last[seq_id].data(), 1, (size_t) n_embd_dec });
+                    common_batch_add(batch, batch_in.token[beg], pending_pos, { seq_id }, /*logits=*/ false);
+                    std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd_dec,
+                                pending_g_last[seq_id].data(), row_bytes);
                }
            }

            for (int32_t k = beg; k < end; ++k) {
-                const int32_t idx = batch.add(batch_in.tokens[k + 1].id, batch_in.tokens[k].pos[0], seq_id, /*output=*/ false);
-                batch.set_embd(idx, { g_embd + (size_t) k * n_embd_dec, 1, (size_t) n_embd_dec });
+                common_batch_add(batch, batch_in.token[k + 1], batch_in.pos[k], { seq_id }, /*logits=*/ false);
+                std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd_dec,
+                            g_embd + (size_t) k * n_embd_dec, row_bytes);
            }

            // refresh deferred state
            const int32_t n_rows = end - beg + 1;
-            verify_pos_first[seq_id] = batch_in.tokens[beg].pos[0];
-            pending_pos_last[seq_id] = batch_in.tokens[end].pos[0];
+            verify_pos_first[seq_id] = batch_in.pos[beg];
+            pending_pos_last[seq_id] = batch_in.pos[end];
            verify_g_rows[seq_id]    = n_rows;
            verify_g[seq_id].resize((size_t) n_rows * n_embd_dec, 0.0f);
            std::memcpy(verify_g[seq_id].data(),       g_embd + (size_t) beg * n_embd_dec, row_bytes * n_rows);
            std::memcpy(pending_g_last[seq_id].data(), g_embd + (size_t) end * n_embd_dec, row_bytes);
        }

-        if (batch.size() > 0) {
-            const int32_t rc = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch.get());
+        if (batch.n_tokens > 0) {
+            const int32_t rc = llama_decode(ctx_dft, batch);
            if (rc != 0) {
-                SPC_ERR("llama_process(ctx_dft) failed rc=%d (n_tokens=%d, ubatch_pos[0]=%d)\n",
-                        rc, (int) batch.size(), (int) batch_in.tokens[0].pos[0]);
+                SPC_ERR("llama_decode(ctx_dft) failed rc=%d (n_tokens=%d, ubatch_pos[0]=%d)\n",
+                        rc, (int) batch.n_tokens, (int) batch_in.pos[0]);
                return false;
            }
        }
@@ -738,12 +719,14 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
    void draft(common_speculative_draft_params_vec & dparams) override {
        auto & ctx_dft = params.ctx_dft;

-        batch.clear();
+        common_batch_clear(batch);

        // keep track of which sequences are still drafting
        int n_drafting = 0;
        std::vector<bool> drafting(n_seq);

+        const size_t row_bytes = (size_t) n_embd_dec * sizeof(float);
+
        // Complete the deferred boundary pair (dp.id_last, pending_g_last) at memory
        // pos pending_pos_last. dp.id_last is target's freshest sample (= corrected
        // token after verify, or first generated token after prefill), matching the
@@ -764,17 +747,19 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {

            llama_memory_seq_rm(llama_get_memory(ctx_dft), seq_id, pending_pos_last[seq_id], -1);

-            const int32_t idx = batch.add(dp.id_last, pending_pos_last[seq_id], seq_id, true);
-            batch.set_embd(idx, { pending_g_last[seq_id].data(), 1, (size_t) n_embd_dec });
+            common_batch_add(batch, dp.id_last, pending_pos_last[seq_id], { seq_id }, true);
+            std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd_dec,
+                        pending_g_last[seq_id].data(),
+                        row_bytes);
        }

-        if (batch.size() == 0) {
+        if (batch.n_tokens == 0) {
            return;
        }

-        int ret = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch.get());
+        int ret = llama_decode(ctx_dft, batch);
        if (ret != 0) {
-            SPC_ERR("llama_process returned %d\n", ret);
+            SPC_ERR("llama_decode returned %d\n", ret);
            return;
        }

@@ -783,7 +768,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
        while (n_drafting > 0) {
            int i_batch = 0;

-            batch.clear();
+            common_batch_clear(batch);

            for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
                if (!drafting[seq_id]) {
@@ -829,17 +814,17 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
                    continue;
                }

-                const int32_t idx = batch.add(id, pending_pos_last[seq_id] + (i + 1), seq_id, true);
-                batch.set_embd(idx, { prenorm, 1, (size_t) n_embd_dec });
+                common_batch_add(batch, id, pending_pos_last[seq_id] + (i + 1), { seq_id }, true);
+                std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd_dec, prenorm, row_bytes);
            }

-            if (batch.size() == 0) {
+            if (batch.n_tokens == 0) {
                break;
            }

-            ret = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch.get());
+            ret = llama_decode(ctx_dft, batch);
            if (ret != 0) {
-                SPC_ERR("llama_process[%d] returned %d\n", i, ret);
+                SPC_ERR("llama_decode[%d] returned %d\n", i, ret);
                break;
            }

@@ -923,10 +908,8 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
 struct common_speculative_impl_draft_dflash : public common_speculative_impl {
    common_params_speculative_draft params;

-    common_batch batch;        // noise tokens
-    common_batch batch_inject; // target features for KV cache injection
-
-    std::vector<float> features_buf; // [n_chunk, n_embd_enc] gathered target features
+    llama_batch batch;        // noise tokens
+    llama_batch batch_inject; // target features for KV cache injection

    std::vector<common_sampler_ptr> smpls;

@@ -1022,11 +1005,15 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
        }
        this->n_max = this->params.n_max;

-        batch        = common_batch(ctx_dft);
-        batch_inject = common_batch(ctx_dft);
+        batch        = llama_batch_init(llama_n_batch(ctx_dft), 0,          n_seq);
+        batch_inject = llama_batch_init(llama_n_ubatch(ctx_dft), n_embd_enc, n_seq);

-        // embd batches on an M-RoPE draft carry 4 position rows per token
+        // embd batches on an M-RoPE draft need 4 position rows per token
        is_mrope = llama_model_rope_type(model_dft) == LLAMA_ROPE_TYPE_MROPE;
+        if (is_mrope) {
+            free(batch_inject.pos);
+            batch_inject.pos = (llama_pos *) malloc(sizeof(llama_pos) * 4 * llama_n_batch(ctx_dft));
+        }

        smpls.resize(n_seq);
        for (auto & s : smpls) {
@@ -1075,6 +1062,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
            llama_sampler_free(backend_chains[seq_id]);
        }
        backend_chains.clear();
+
+        llama_batch_free(batch);
+        llama_batch_free(batch_inject);
    }

    void begin(llama_seq_id seq_id, const llama_tokens & prompt) override {
@@ -1095,8 +1085,8 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
        }
    }

-    bool process(const common_batch & batch_in) override {
-        if (batch_in.size() <= 0) {
+    bool process(const llama_batch & batch_in) override {
+        if (batch_in.n_tokens <= 0) {
            return true;
        }

@@ -1104,19 +1094,20 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
        // produce the target-layer features used to seed the draft KV cache, so
        // embeddings are injected too, except the pinned ones skipped below.
        // TODO: revisit after https://github.com/ggml-org/llama.cpp/pull/24669 is merged
-        const bool has_tokens     = batch_in.has_token();
-        const bool has_embeddings = batch_in.has_embd();
+        const bool has_tokens     = batch_in.token != nullptr;
+        const bool has_embeddings = batch_in.embd  != nullptr;
        if (has_tokens == has_embeddings) {
            return true;
        }

-        const int32_t n_tokens = batch_in.size();
+        const int32_t n_tokens = batch_in.n_tokens;

        // per-seq inclusive batch range (assumes each seq's tokens are contiguous in the batch)
        std::vector<int32_t> i_batch_beg(n_seq, -1);
        std::vector<int32_t> i_batch_end(n_seq, -1);
        for (int32_t k = 0; k < n_tokens; ++k) {
-            const llama_seq_id seq_id = batch_in.tokens[k].seq_id;
+            GGML_ASSERT(batch_in.n_seq_id[k] == 1);
+            const llama_seq_id seq_id = batch_in.seq_id[k][0];
            if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) {
                continue;
            }
@@ -1139,7 +1130,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {

            // an M-RoPE image pins all its rows to one position, so a windowed draft
            // cache cannot free cells for it - skip it, the draft can jump over the gap
-            const bool pos_pinned = batch_in.tokens[i_batch_beg[seq_id]].pos[0] == batch_in.tokens[i_batch_end[seq_id]].pos[0];
+            const bool pos_pinned = batch_in.pos[i_batch_beg[seq_id]] == batch_in.pos[i_batch_end[seq_id]];
            if (has_embeddings && n_rows > 1 && pos_pinned) {
                continue;
            }
@@ -1149,28 +1140,34 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {

                // gather target features per extract layer; the fused decode encodes and
                // injects them into the K/V cache at the target positions
-                features_buf.resize((size_t) n_chunk * n_embd_enc);
+                batch_inject.n_tokens = n_chunk;
                for (uint32_t k = 0; k < target_layer_ids_n; ++k) {
                    const float * layer = llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k]);
                    if (!layer) {
                        GGML_ABORT("DFlash: target layer %d input not extracted.", target_layer_ids[k]);
                    }
                    for (int32_t i = 0; i < n_chunk; ++i) {
-                        float       * dst = features_buf.data() + (size_t) i * n_embd_enc + k * (size_t) n_embd_tgt;
+                        float       * dst = batch_inject.embd + (size_t) i * n_embd_enc + k * (size_t) n_embd_tgt;
                        const float * src = layer + (size_t) (i_batch_beg[seq_id] + offset + i) * n_embd_tgt;
                        std::memcpy(dst, src, (size_t) n_embd_tgt * sizeof(float));
                    }
                }

-                batch_inject.clear();
                for (int32_t i = 0; i < n_chunk; ++i) {
-                    const llama_pos p = batch_in.tokens[i_batch_beg[seq_id] + offset + i].pos[0];
-                    const llama_pos pos_arr[4] = { p, p, p, 0 };
-                    batch_inject.add_embd({ features_buf.data() + (size_t) i * n_embd_enc, 1, (size_t) n_embd_enc }, pos_arr, seq_id, false);
+                    const llama_pos p = batch_in.pos[i_batch_beg[seq_id] + offset + i];
+                    batch_inject.pos[i] = p;
+                    if (is_mrope) {
+                        batch_inject.pos[1 * n_chunk + i] = p;
+                        batch_inject.pos[2 * n_chunk + i] = p;
+                        batch_inject.pos[3 * n_chunk + i] = 0;
+                    }
+                    batch_inject.n_seq_id[i]  = 1;
+                    batch_inject.seq_id[i][0] = seq_id;
+                    batch_inject.logits[i]    = false;
                }
-                const int32_t rc = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch_inject.get());
+                const int32_t rc = llama_decode(ctx_dft, batch_inject);
                if (rc != 0) {
-                    LOG_ERR("%s: llama_process(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
+                    LOG_ERR("%s: llama_decode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
                            __func__, rc, (int) n_chunk, (int) offset);
                    return false;
                }
@@ -1183,7 +1180,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
    void draft(common_speculative_draft_params_vec & dparams) override {
        auto & ctx_dft = params.ctx_dft;

-        batch.clear();
+        common_batch_clear(batch);

        // build one batch holding every drafting sequence's noise block into a single decode)
        // record where each block starts and its size
@@ -1203,21 +1200,21 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
            const int32_t n_draft = params.n_max;

            const int32_t n_block_tokens = n_draft + (is_dspark && sample_from_anchor ? 0 : 1);
-            i_block_beg[seq_id] = batch.size();
+            i_block_beg[seq_id] = batch.n_tokens;
            n_block    [seq_id] = n_block_tokens;
            for (int32_t i = 0; i < n_block_tokens; ++i) {
-                batch.add(i == 0 ? dp.id_last : mask_token_id, n + i, seq_id, !is_dflash2);
+                common_batch_add(batch, i == 0 ? dp.id_last : mask_token_id, n + i, { seq_id }, !is_dflash2);
            }
        }

-        if (batch.size() == 0) {
+        if (batch.n_tokens == 0) {
            return;
        }

        // decode all sequence's noise block in a single batch
-        int ret = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch.get());
+        int ret = llama_decode(ctx_dft, batch);
        if (ret != 0) {
-            LOG_WRN("%s: llama_process returned %d\n", __func__, ret);
+            LOG_WRN("%s: llama_decode returned %d\n", __func__, ret);
            return;
        }

@@ -1331,7 +1328,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
 struct common_speculative_impl_draft_mtp : public common_speculative_impl {
    common_params_speculative_draft params; // reuses the draft-model params slot (ctx_tgt/ctx_dft)

-    common_batch batch;
+    llama_batch batch;

    std::vector<common_sampler_ptr> smpls;

@@ -1387,7 +1384,11 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
                ctx_dft ? "yes" : "no",
                common_speculative_get_devices_str(this->params.devices).c_str());

-        batch = common_batch(ctx_dft);
+        const int32_t n_b = (int32_t) llama_n_batch(ctx_dft);
+        batch = llama_batch_init(/*n_tokens=*/ n_b, /*embd=*/ n_embd, /*n_seq_max=*/ 1);
+        // llama_batch_init allocates only one of token/embd; MTP needs both.
+        // TODO: fix, how to call without malloc
+        batch.token = (llama_token *) malloc(sizeof(llama_token) * n_b);

        smpls.resize(n_seq);
        for (auto & s : smpls) {
@@ -1452,6 +1453,12 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
            llama_sampler_free(backend_chains[seq_id]);
        }
        backend_chains.clear();
+
+        if (batch.token != nullptr) {
+            free(batch.token);
+            batch.token = nullptr;
+        }
+        llama_batch_free(batch);
    }

    void begin(llama_seq_id seq_id, const llama_tokens & prompt) override {
@@ -1466,23 +1473,23 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
        if (pos_max < N - 1 && !is_mem_shared) {
            SPC_WRN("ctx_dft pos_max=%d < N-1=%d - "
                    "process() hook may not have run on every prefill ubatch "
-                    "(need_embd / output flag on every prompt position?). "
+                    "(need_embd / logits=1 on every prompt position?). "
                    "Drafts may degrade.\n",
                    (int) pos_max, N - 1);
        }
    }

-    bool process(const common_batch & batch_in) override {
-        if (batch_in.size() <= 0) {
+    bool process(const llama_batch & batch_in) override {
+        if (batch_in.n_tokens <= 0) {
            return true;
        }

        // TODO: how to make it work with vision tokens?
-        if (!batch_in.has_token() || batch_in.has_embd()) {
+        if (batch_in.token == nullptr || batch_in.embd != nullptr) {
            return true;
        }

-        const int32_t n_tokens = batch_in.size();
+        const int32_t n_tokens = batch_in.n_tokens;

        // remember the first and last batch index for each sequence
        std::fill(i_batch_beg.begin(), i_batch_beg.end(), -1);
@@ -1490,7 +1497,9 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {

        for (int k = 0; k < n_tokens; ++k) {
            for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
-                if (batch_in.tokens[k].seq_id == seq_id) {
+                GGML_ASSERT(batch_in.n_seq_id[k] == 1);
+
+                if (batch_in.seq_id[k][0] == seq_id) {
                    i_batch_end[seq_id] = k;
                    if (i_batch_beg[seq_id] < 0) {
                        i_batch_beg[seq_id] = k;
@@ -1506,26 +1515,33 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {

        // if kv is shared with target (e.g Gemma4), then we can skip this catch-up decode
        if (!is_mem_shared) {
-            batch.clear();
+            common_batch_clear(batch);

-            // pair each token with the tgt embedding shifted right by one position, and
-            // the first token of each sequence with the pending embedding from a previous run
+            for (int k = 0; k < n_tokens; ++k) {
+                common_batch_add(batch, batch_in.token[k], batch_in.pos[k], { batch_in.seq_id[k][0] }, 0);
+            }
+
+            // shift the tgt embeddings to the right by one position
            // assumes that the tokens in the batch are sequential for each sequence
            // i.e. we cannot have seq_id like this: [0, 0, 0, 1, 1, 0, 1, 1]
            //                                                       ^--- this is a problem
            // TODO:this is generally true, but would be nice to assert it
-            const float * h_tgt = llama_get_embeddings_nextn(ctx_tgt);
+            {
+                const float * h_tgt = llama_get_embeddings_nextn(ctx_tgt);
+                std::memcpy(batch.embd + (size_t) 1 * n_embd, h_tgt, row_bytes * (n_tokens-1));
+            }

-            for (int k = 0; k < n_tokens; ++k) {
-                const llama_seq_id seq_id = batch_in.tokens[k].seq_id;
+            // fill the pending embeddings from a previous run
+            auto set_h = [&](int idx, const float * h_row) {
+                std::memcpy(batch.embd + (size_t) idx * n_embd, h_row, row_bytes);
+            };

-                const int32_t idx = batch.add(batch_in.tokens[k].id, batch_in.tokens[k].pos[0], seq_id, false);
+            for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
+                if (i_batch_beg[seq_id] < 0) {
+                    continue;
+                }

-                const float * h_row = k == i_batch_beg[seq_id]
-                    ? pending_h[seq_id].data()
-                    : h_tgt + (size_t) (k - 1) * n_embd;
-
-                batch.set_embd(idx, { h_row, 1, (size_t) n_embd });
+                set_h(i_batch_beg[seq_id], pending_h[seq_id].data());
            }

            auto * mem_dft = llama_get_memory(ctx_dft);
@@ -1538,15 +1554,15 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
                        if (i_batch_beg[seq_id] < 0) {
                            continue;
                        }
-                        llama_memory_seq_rm(mem_dft, seq_id, batch_in.tokens[i_batch_beg[seq_id]].pos[0], -1);
+                        llama_memory_seq_rm(mem_dft, seq_id, batch_in.pos[i_batch_beg[seq_id]], -1);
                    }
                    llama_set_nextn_layer_offset(ctx_dft, head);
                }

-                const int32_t rc = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch.get());
+                const int32_t rc = llama_decode(ctx_dft, batch);
                if (rc != 0) {
-                    SPC_ERR("llama_process(ctx_dft) head=%d failed rc=%d (pos=%d)\n",
-                            head, (int) rc, (int) batch_in.tokens[0].pos[0]);
+                    SPC_ERR("llama_decode(ctx_dft) head=%d failed rc=%d (pos=%d)\n",
+                            head, (int) rc, (int) batch_in.pos[0]);
                    ok = false;
                    break;
                }
@@ -1584,12 +1600,14 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
    void draft(common_speculative_draft_params_vec & dparams) override {
        auto & ctx_dft = params.ctx_dft;

-        batch.clear();
+        common_batch_clear(batch);

        // keep track of which sequences are still drafting
        int n_drafting = 0;
        std::vector<bool> drafting(n_seq);

+        const size_t row_bytes = (size_t) n_embd * sizeof(float);
+
        for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
            auto & dp = dparams[seq_id];

@@ -1601,10 +1619,10 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
            drafting[seq_id] = true;
            common_sampler_reset(smpls[seq_id].get());

-            const int32_t idx = batch.add(dp.id_last, dp.pos0, seq_id, true);
-            batch.set_embd(idx, { pending_h[seq_id].data(), 1, (size_t) n_embd });
+            common_batch_add(batch, dp.id_last, dp.pos0, { seq_id }, true);
+            std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, pending_h[seq_id].data(), row_bytes);

-            i_last[seq_id] = idx;
+            i_last[seq_id] = batch.n_tokens - 1;

            if (chain_heads) {
                chain_h[seq_id].assign(pending_h[seq_id].begin(), pending_h[seq_id].end());
@@ -1630,16 +1648,16 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
                llama_set_nextn_layer_offset(ctx_dft, i);
            }

-            int ret = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch.get());
+            int ret = llama_decode(ctx_dft, batch);
            if (ret != 0) {
-                SPC_ERR("llama_process[%d] returned %d\n", i, ret);
+                SPC_ERR("llama_decode[%d] returned %d\n", i, ret);
                break;
            }

            // rebuild the batch for the next step: the growing-KV paths re-add only the
            // new token (the KV already holds the prefix), while chained heads re-add the
            // whole prefix at the next head. dropped sequences are simply not re-added.
-            batch.clear();
+            common_batch_clear(batch);

            for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
                if (!drafting[seq_id]) {
@@ -1690,24 +1708,24 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
                    const int n_rows = (int) result.size() + 1; // id_last + tokens drafted so far
                    for (int t = 0; t < n_rows; ++t) {
                        const llama_token tok = (t == 0) ? dp.id_last : result[t - 1];
-                        const int32_t idx = batch.add(tok, dp.pos0 + t, seq_id, t == n_rows - 1);
-                        batch.set_embd(idx, { chain_h[seq_id].data() + (size_t) t * n_embd, 1, (size_t) n_embd });
-                        i_last[seq_id] = idx;
+                        common_batch_add(batch, tok, dp.pos0 + t, { seq_id }, t == n_rows - 1);
+                        std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd,
+                                    chain_h[seq_id].data() + (size_t) t * n_embd, row_bytes);
                    }
                } else if (is_mem_shared) {
                    // note: with shared memory (e.g. Gemma4 assistants) we use the same position for all draft tokens
                    // ref: https://github.com/huggingface/transformers/blob/effde20942e3f82a1b97449f60b3a48c5ff96145/docs/source/en/model_doc/gemma4_assistant.md?plain=1#L36-L37
-                    const int32_t idx = batch.add(id, dp.pos0, seq_id, true);
-                    batch.set_embd(idx, { h_row, 1, (size_t) n_embd });
-                    i_last[seq_id] = idx;
+                    common_batch_add(batch, id, dp.pos0, { seq_id }, true);
+                    std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, h_row, row_bytes);
                } else {
-                    const int32_t idx = batch.add(id, dp.pos0 + i + 1, seq_id, true);
-                    batch.set_embd(idx, { h_row, 1, (size_t) n_embd });
-                    i_last[seq_id] = idx;
+                    common_batch_add(batch, id, dp.pos0 + i + 1, { seq_id }, true);
+                    std::memcpy(batch.embd + (size_t) (batch.n_tokens - 1) * n_embd, h_row, row_bytes);
                }
+
+                i_last[seq_id] = batch.n_tokens - 1;
            }

-            if (batch.size() == 0) {
+            if (batch.n_tokens == 0) {
                break;
            }

@@ -1769,7 +1787,7 @@ struct common_speculative_impl_ngram_simple : public common_speculative_impl {
        // noop
    }

-    bool process(const common_batch & /*batch*/) override {
+    bool process(const llama_batch & /*batch*/) override {
        // TODO: implement
        return true;
    }
@@ -1817,7 +1835,7 @@ struct common_speculative_impl_ngram_map_k : public common_speculative_impl {
        common_ngram_map_begin(config[seq_id], prompt);
    }

-    bool process(const common_batch & /*batch*/) override {
+    bool process(const llama_batch & /*batch*/) override {
        // TODO: implement
        return true;
    }
@@ -1975,7 +1993,7 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl {
        sinfo.n_draft_last = result.size();
    }

-    bool process(const common_batch & /*batch*/) override {
+    bool process(const llama_batch & /*batch*/) override {
        // TODO: implement
        return true;
    }
@@ -2137,7 +2155,7 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl {
        }
    }

-    bool process(const common_batch & /*batch*/) override {
+    bool process(const llama_batch & /*batch*/) override {
        // TODO: implement
        return true;
    }
@@ -2539,7 +2557,7 @@ common_speculative_init_result::common_speculative_init_result(
        model_path = params.speculative.draft.mparams.path;
        LOG_INF("%s: loading draft model '%s'\n", __func__, model_path.c_str());

-        llama_model * model_dft = llama_model_load_from_file(model_path.c_str(), mparams);
+        llama_model * model_dft = llama_model_load_from_file(params.model.path.c_str(), mparams);
        if (model_dft == NULL) {
            LOG_ERR("%s: failed to load draft model, '%s'\n", __func__, model_path.c_str());
            return;
@@ -2770,7 +2788,7 @@ void common_speculative_begin(common_speculative * spec, llama_seq_id seq_id, co
    }
 }

-bool common_speculative_process(common_speculative * spec, const common_batch & batch) {
+bool common_speculative_process(common_speculative * spec, const llama_batch & batch) {
    bool result = true;

    if (spec == nullptr) {
@@ -77,7 +77,7 @@ common_speculative_draft_params & common_speculative_get_draft_params(common_spe
 void common_speculative_begin(common_speculative * spec, llama_seq_id seq_id, const llama_tokens & prompt);

 // process the batch and update the internal state of the speculative context
-bool common_speculative_process(common_speculative * spec, const common_batch & batch);
+bool common_speculative_process(common_speculative * spec, const llama_batch & batch);

 // generate drafts for the sequences specified with `common_speculative_get_draft_params`
 void common_speculative_draft(common_speculative * spec);
@@ -106,7 +106,6 @@ TEXT_MODEL_MAP: dict[str, str] = {
    "Glm4MoeLiteForCausalLM": "glm",
    "Glm4vForConditionalGeneration": "glm",
    "Glm4vMoeForConditionalGeneration": "glm",
-    "Glm5NextForConditionalGeneration": "glm",
    "GlmForCausalLM": "chatglm",
    "GlmMoeDsaForCausalLM": "glm",
    "GlmOcrForConditionalGeneration": "glm",
@@ -306,7 +305,6 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
    "Glm4vForConditionalGeneration": "qwen3vl",
    "BailingMoeV3VLForConditionalGeneration": "bailingmoe3",
    "Glm4vMoeForConditionalGeneration": "qwen3vl",
-    "Glm5NextForConditionalGeneration": "qwen3vl",
    "Glm5vForConditionalGeneration": "kimivl",
    "GlmOcrForConditionalGeneration": "qwen3vl",
    "GlmasrModel": "ultravox",
@@ -234,7 +234,7 @@ class ModelBase:

        prefix = "model" if not self.is_mistral_format else "consolidated"
        part_names: list[str] = ModelBase.get_model_part_names(self.dir_model, prefix, ".safetensors")
-        is_safetensors: bool = len(part_names) > 0 or (not self.is_mistral_format and (self.dir_model / "model.safetensors.index.json").is_file())
+        is_safetensors: bool = len(part_names) > 0
        if not is_safetensors:
            part_names = ModelBase.get_model_part_names(self.dir_model, "pytorch_model", ".bin")

@@ -2326,12 +2326,6 @@ class TextModel(ModelBase):
                raise NotImplementedError("Only MEAN, CLS, and LAST pooling types supported")
            self.gguf_writer.add_pooling_type(pooling_type)

-        # pooling before a classification head (e.g. ModernBertForSequenceClassification)
-        if (classifier_pooling := self.hparams.get("classifier_pooling")) is not None:
-            if classifier_pooling not in ("cls", "mean"):
-                raise NotImplementedError(f"Unsupported classifier_pooling: {classifier_pooling}")
-            self.gguf_writer.add_classifier_pooling_type(mode_mapping[classifier_pooling])
-
    def _set_vocab_glmedge(self):
        from transformers import AutoTokenizer
        tokenizer = AutoTokenizer.from_pretrained(self.dir_model)
@@ -2566,11 +2560,6 @@ class TextModel(ModelBase):

        self.gguf_writer.add_add_space_prefix(False)

-        if (add_bos := tokenizer_config.get("add_bos_token")) is not None:
-            self.gguf_writer.add_add_bos_token(add_bos)
-        if (add_eos := tokenizer_config.get("add_eos_token")) is not None:
-            self.gguf_writer.add_add_eos_token(add_eos)
-

 class MmprojModel(ModelBase):
    model_type = ModelType.MMPROJ
@@ -402,197 +402,3 @@ class SolarOpenModel(Glm4MoeModel):
        special_vocab._set_special_token("unk", tokenizer.get_added_vocab()["<unk>"])  # ty: ignore[unresolved-attribute]
        special_vocab._set_special_token("bos", tokenizer.get_added_vocab()["<|startoftext|>"])  # ty: ignore[unresolved-attribute]
        special_vocab.add_to_gguf(self.gguf_writer)
-
-
-@ModelBase.register("Glm5NextForConditionalGeneration")
-@ModelBase.example("zai-org/GLM-5.3-Flash")
-class Glm5NextModel(TextModel):
-
-    model_arch = gguf.MODEL_ARCH.GLM5_NEXT
-    supports_mtp_export = True
-
-    _experts: list[dict[str, Tensor]] | None = None
-    _n_main_layers: int | None = None
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        self.n_nextn_layers = self.hparams.get("num_nextn_predict_layers", 0)
-        self.skip_mtp = self.no_mtp or self.n_nextn_layers == 0
-
-        if not self.skip_mtp:
-            self.block_count += self.n_nextn_layers
-            self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
-
-        self.hparams.pop("head_dim", None)
-
-    def set_vocab(self):
-        # requires transformers >= 5, tokpre hash-resolves to glm4
-        return self._set_vocab_glm()
-
-    def index_tensors(self, remote_hf_model_id: str | None = None):
-        hp = self.hparams.get("text_config", self.hparams)
-        type(self)._n_main_layers = hp["num_hidden_layers"]
-        return super().index_tensors(remote_hf_model_id=remote_hf_model_id)
-
-    @classmethod
-    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
-        if (titem := super().filter_tensors(item)) is None:
-            return None
-        name, gen = titem
-
-        assert cls._n_main_layers is not None
-        m = re.match(r"model\.layers\.(\d+)\.", name)
-        is_mtp = m is not None and int(m.group(1)) >= cls._n_main_layers
-
-        if is_mtp and cls.no_mtp:
-            return None
-        if cls.mtp_only and not is_mtp and name not in (
-            "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",
-        ):
-            return None
-
-        return name, gen
-
-    def set_gguf_parameters(self):
-        super().set_gguf_parameters()
-        hp = self.hparams
-
-        layer_types = hp["layer_types"]
-        n_kv_heads = [0 if t == "linear_attention" else 1 for t in layer_types]
-        assert len(n_kv_heads) == hp["num_hidden_layers"]
-        # Pad to block_count
-        n_kv_heads += [1] * (self.block_count - len(n_kv_heads))
-        self.gguf_writer.add_head_count_kv(n_kv_heads)
-        self.gguf_writer.add_vocab_size(hp["vocab_size"])
-        self.gguf_writer.add_layer_norm_eps(1e-6)
-
-        if not self.skip_mtp:
-            self.gguf_writer.add_nextn_predict_layers(self.n_nextn_layers)
-
-        # KDA
-        lin = hp["linear_attn_config"]
-        assert lin["num_heads"] == hp["num_attention_heads"]
-        self.gguf_writer.add_ssm_conv_kernel(lin["short_conv_kernel_size"])
-        self.gguf_writer.add_kda_head_dim(lin["head_dim"])
-        if (lb := lin.get("gate_lower_bound")) is not None:
-            self.gguf_writer.add_kda_gate_lower_bound(lb)
-
-        # MLA (nope only)
-        assert hp.get("mla_use_nope") and hp["qk_rope_head_dim"] == 0, "expected nope-only MLA"
-        kv_lora_rank = hp["kv_lora_rank"]
-        qk_rope = hp["qk_rope_head_dim"]
-        self.gguf_writer.add_q_lora_rank(hp["q_lora_rank"])
-        self.gguf_writer.add_kv_lora_rank(kv_lora_rank)
-        self.gguf_writer.add_rope_dimension_count(qk_rope)
-        self.gguf_writer.add_key_length(kv_lora_rank + qk_rope)
-        self.gguf_writer.add_value_length(kv_lora_rank)
-        self.gguf_writer.add_key_length_mla(hp["qk_nope_head_dim"] + qk_rope)
-        self.gguf_writer.add_value_length_mla(hp["v_head_dim"])
-
-        # DSA indexer with k-pool compression
-        self.gguf_writer.add_indexer_head_count(hp["index_n_heads"])
-        self.gguf_writer.add_indexer_key_length(hp["index_head_dim"])
-        self.gguf_writer.add_indexer_top_k(hp["index_topk"])
-        self.gguf_writer.add_indexer_kpool(hp["index_kpool"])
-        self.gguf_writer.add_indexer_kpool_select_tail(hp.get("index_kpool_always_select_tail", True))
-        if (indexer_types := hp.get("indexer_types")) is not None:
-            self.gguf_writer.add_indexer_types([t == "full" for t in indexer_types])
-
-        # mHC
-        assert hp.get("mhc", True)
-        self.gguf_writer.add_hyper_connection_count(hp["hc_mult"])
-        self.gguf_writer.add_hyper_connection_sinkhorn_iterations(hp["hc_sinkhorn_iters"])
-        self.gguf_writer.add_hyper_connection_epsilon(hp["hc_eps"])
-
-        # MoE
-        self.gguf_writer.add_leading_dense_block_count(hp["first_k_dense_replace"])
-        self.gguf_writer.add_expert_feed_forward_length(hp["moe_intermediate_size"])
-        self.gguf_writer.add_expert_shared_count(hp["n_shared_experts"])
-        self.gguf_writer.add_expert_weights_scale(hp["routed_scaling_factor"])
-        self.gguf_writer.add_expert_weights_norm(hp["norm_topk_prob"])
-        if (limit := hp.get("swiglu_limit")) is not None:
-            self.gguf_writer.add_swiglu_clamp_exp([limit] * self.block_count)
-            self.gguf_writer.add_swiglu_clamp_shexp([limit] * self.block_count)
-
-    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        if name == "lm_head.weight" and self.hparams.get("tie_word_embeddings", False):
-            return
-
-        # routed experts
-        if ".mlp.experts." in name:
-            n_experts = self.hparams["n_routed_experts"]
-            assert bid is not None
-            if self._experts is None:
-                self._experts = [{} for _ in range(self.block_count)]
-            self._experts[bid][name] = data_torch
-            if len(self._experts[bid]) < n_experts * 3:
-                return
-            for w_name in ("down_proj", "gate_proj", "up_proj"):
-                datas: list[Tensor] = []
-                for xid in range(n_experts):
-                    ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
-                    datas.append(self._experts[bid].pop(ename))
-                merged = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
-                yield from super().modify_tensors(torch.stack(datas, dim=0), merged, bid)
-            return
-
-        # MLA absorption
-        if name.endswith("kv_b_proj.weight"):
-            n_head = self.hparams["num_attention_heads"]
-            v_head_dim = self.hparams["v_head_dim"]
-            qk_nope_head_dim = self.hparams["qk_nope_head_dim"]
-            assert data_torch.shape[0] == n_head * (v_head_dim + qk_nope_head_dim)
-            kv_b = data_torch.view(n_head, v_head_dim + qk_nope_head_dim, data_torch.shape[-1])
-            k_b, v_b = torch.split(kv_b, [qk_nope_head_dim, v_head_dim], dim=1)
-            yield from super().modify_tensors(k_b.transpose(1, 2), name.replace("kv_b_proj", "k_b_proj"), bid)
-            yield from super().modify_tensors(v_b, name.replace("kv_b_proj", "v_b_proj"), bid)
-            return
-
-        # KDA conv1d
-        if name.endswith((".q_conv1d.weight", ".k_conv1d.weight", ".v_conv1d.weight")):
-            if data_torch.ndim == 3:
-                d_inner, _, d_conv = data_torch.shape
-            elif data_torch.ndim == 2:
-                d_inner, d_conv = data_torch.shape
-            else:
-                raise ValueError(f"unexpected conv1d rank {data_torch.ndim} for {name}")
-            data_torch = data_torch.reshape(1, d_inner, 1, d_conv)
-
-        if name.endswith(".A_log"):
-            n_head = self.hparams["num_attention_heads"]
-            data_torch = -torch.exp(data_torch.float().flatten()[:n_head])
-
-        if name.endswith(".dt_bias"):
-            name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias"
-
-        if re.search(r"\.(hc_(?:attn|ffn)_(?:fn|base|scale)|index_kpool_compress_(?:ape|gate))$", name):
-            yield self.map_tensor_name(name) + ".weight", data_torch
-            return
-
-        yield from super().modify_tensors(data_torch, name, bid)
-
-    def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool:
-        # keep the small mHC / gating parameters exact
-        exact_keys = ("hc_attn_", "hc_ffn_", "indexer_compressor_", "ssm_a", "ssm_dt", "exp_probs_b")
-        if new_name.startswith(("blk.", "output_hc")) and any(k in new_name for k in exact_keys):
-            return gguf.GGMLQuantizationType.F32
-        return super().tensor_force_quant(name, new_name, bid, n_dims)
-
-    def prepare_metadata(self, vocab_only: bool):
-        from_dir = self.fname_out.is_dir()
-        super().prepare_metadata(vocab_only=vocab_only)
-        if not self.mtp_only or not from_dir:
-            return
-        output_type: str = self.ftype.name.partition("_")[2]
-        fname_default: str = gguf.naming_convention(
-            self.metadata.name, self.metadata.basename, self.metadata.finetune,
-            self.metadata.version, size_label=None, output_type=output_type, model_type=None)
-        self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"
-
-    def prepare_tensors(self):
-        super().prepare_tensors()
-        if self._experts is not None:
-            leftover = [k for d in self._experts for k in d.keys()]
-            if leftover:
-                raise ValueError(f"Unprocessed experts: {leftover}")
@@ -469,21 +469,6 @@ class _LinearAttentionVReorderBase(Qwen3NextModel):
        shape = list(tensor.shape)
        if dim < 0:
            dim += len(shape)
-
-        # LoRA tensors (W ≈ B @ A) cannot reshape their row dimension.
-        # Instead, build a permutation index and apply it to A (column reorder) or B (row reorder) directly.
-        if hasattr(tensor, 'get_lora_A_B'):
-            n = shape[dim]
-            idx = torch.arange(n).reshape(num_k_heads, num_v_per_k, head_dim)
-            idx = idx.permute(1, 0, 2).contiguous().reshape(n)
-            lora_A, lora_B = tensor.get_lora_A_B()  # ty: ignore[call-non-callable]
-            if dim == len(shape) - 1:
-                return type(tensor)(lora_A[:, idx], lora_B)
-            elif dim == 0:
-                return type(tensor)(lora_A, lora_B[idx])
-            else:
-                raise NotImplementedError(f"_reorder_v_heads on dim={dim} not supported for LoRA tensors")
-
        new_shape = shape[:dim] + [num_k_heads, num_v_per_k, head_dim] + shape[dim + 1:]
        tensor = tensor.reshape(*new_shape)
        perm = list(range(len(new_shape)))
@@ -701,12 +686,6 @@ class DFlashModel(Qwen3Model):
        super().set_gguf_parameters()

        dflash_config = self.hparams.get("dflash_config", {})
-        if (partial_rotary_factor := self.rope_parameters.get("partial_rotary_factor")) is not None:
-            head_dim = self.hparams.get("head_dim") or self.hparams["hidden_size"] // self.hparams["num_attention_heads"]
-            self.gguf_writer.add_rope_dimension_count(int(head_dim * partial_rotary_factor))
-        if (value_scale := dflash_config.get("attention_value_scale")) is not None:
-            self.gguf_writer.add_attn_value_scale(float(value_scale))
-
        block_size = dflash_config.get("block_size", self.hparams.get("block_size", 16))
        self.gguf_writer.add_block_size(block_size)

@@ -752,67 +731,17 @@ class DFlashModel(Qwen3Model):
        if causal is not None:
            self.gguf_writer.add_causal_attention(bool(causal))

-        # M-RoPE target: the draft ropes on the temporal dim only, so write
-        # degenerate sections [n_rot/2, 0, 0, 0]
+        head_dim = self.hparams.get("head_dim") or self.hparams["hidden_size"] // self.hparams["num_attention_heads"]
        if self._target_uses_mrope():
-            head_dim = self.hparams.get("head_dim") or self.hparams["hidden_size"] // self.hparams["num_attention_heads"]
+            # M-RoPE target: the draft ropes on the temporal dim only, so write degenerate sections [n_rot/2, 0, 0, 0]
            self.gguf_writer.add_rope_dimension_sections([head_dim // 2, 0, 0, 0])
+        elif (prf := self.hparams.get("partial_rotary_factor")) is not None:
+            # partial rotary (e.g. MiMo V2 DFlash): only a fraction of the head dims rotate
+            self.gguf_writer.add_rope_dimension_count(int(head_dim * float(prf)))

-    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
-        yield from super().generate_extra_tensors()
-
-        mask_path = self.dir_model / "mask_embedding.pt"
-        if not mask_path.is_file():
-            return
-
-        mask = torch.load(mask_path, map_location="cpu", weights_only=True)
-        mask_id = self.hparams.get("dflash_config", {}).get("mask_token_id")
-        if mask_id is None or mask["mask_token_id"] != mask_id:
-            raise ValueError("mask_embedding.pt mask_token_id does not match dflash_config")
-        if tuple(mask["embedding"].shape) != (self.hparams["hidden_size"],):
-            raise ValueError("mask_embedding.pt has an unexpected embedding shape")
-        if not 0 <= mask_id < self.hparams["vocab_size"]:
-            raise ValueError("mask_embedding.pt mask_token_id is outside the vocabulary")
-
-        def target_tensor(name: str) -> Tensor:
-            if self.target_model_dir is None:
-                raise ValueError("mask_embedding.pt requires --target-model-dir with the target embeddings and output head")
-            index_path = self.target_model_dir / "model.safetensors.index.json"
-            if index_path.is_file():
-                with open(index_path, encoding="utf-8") as f:
-                    weight_map = json.load(f)["weight_map"]
-                part_names = [weight_map[name]]
-            else:
-                part_names = self.get_model_part_names(self.target_model_dir, "model", ".safetensors")
-
-            for part_name in part_names:
-                with gguf.utility.SafetensorsLocal(self.target_model_dir / part_name) as part:
-                    if name in part:
-                        return LazyTorchTensor.from_local_tensor(part[name])
-            raise ValueError(f"Target tensor {name!r} was not found in safetensors")
-
-        embedding_name = "model.embed_tokens.weight"
-        if embedding_name in self.model_tensors:
-            embeddings = self.model_tensors.pop(embedding_name)()
-        else:
-            embeddings = target_tensor(embedding_name)
-
-        if "model.lm_head.weight" not in self.model_tensors:
-            if self.target_model_dir is None:
-                raise ValueError("mask_embedding.pt requires --target-model-dir to obtain the output head")
-            target_config = ModelBase.load_hparams(self.target_model_dir, False)
-            target_config = {**target_config, **target_config.get("text_config", {})}
-            head_name = embedding_name if target_config.get("tie_word_embeddings", False) else "lm_head.weight"
-            # Keep the output head separate from the patched input embedding table.
-            yield "model.lm_head.weight", target_tensor(head_name)
-
-        embeddings = LazyTorchTensor.to_eager(embeddings).clone()
-        if tuple(embeddings.shape) != (self.hparams["vocab_size"], self.hparams["hidden_size"]):
-            raise ValueError("Target token embedding shape does not match the DFlash draft")
-        # MiMo's target mask row is untrained; the draft provides its own vector.
-        embeddings[mask_id] = mask["embedding"].to(embeddings.dtype)
-        self.hparams["has_embed_tokens"] = True
-        yield embedding_name, embeddings
+        # attention value scale (e.g. MiMo V2 DFlash): applied on the attention output
+        if (v_scale := dflash_config.get("attention_value_scale", self.hparams.get("attention_value_scale"))) is not None:
+            self.gguf_writer.add_attn_value_scale(float(v_scale))

    def _target_uses_mrope(self) -> bool:
        if self.target_model_dir is None:
@@ -823,6 +752,69 @@ class DFlashModel(Qwen3Model):
        rope = cfg.get("rope_parameters") or cfg.get("rope_scaling") or {}
        return "mrope_section" in rope

+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        # The target's embedding row for the mask token may be untrained (all zeros), but the
+        # drafter was trained with a separate learned vector (mask_embedding.pt). Synthesize the
+        # draft's own token_embd from the target's table with that row replaced; a runtime that
+        # embeds mask slots through the target's table would otherwise feed the drafter zeros.
+        # ref: https://huggingface.co/TrevorJS/Xiaomi-MiMo-V2.6-Flash-RL-DFlash
+        dflash_config = self.hparams.get("dflash_config", {})
+        mask_token_id = dflash_config.get("mask_token_id", self.hparams.get("mask_token_id"))
+        mask_emb_file = self.dir_model / "mask_embedding.pt"
+        has_own_embd = any(k.endswith("embed_tokens.weight") for k in self.model_tensors)
+        if (self.target_model_dir is not None and mask_token_id is not None
+                and mask_emb_file.is_file() and not has_own_embd):
+            target_embd = self._load_target_embed_tokens()
+            if target_embd is None:
+                logger.warning("DFlash: could not load the target token_embd to patch the mask embedding")
+            else:
+                embd = target_embd.detach().clone()
+                mask = torch.load(str(mask_emb_file), map_location="cpu", weights_only=True)
+                # the file may wrap the vector in a dict/list rather than being a bare tensor
+                if isinstance(mask, (dict, list, tuple)):
+                    mask = next(iter(mask.values()) if isinstance(mask, dict) else iter(mask))
+                mask = torch.as_tensor(mask).reshape(-1)
+                if mask.numel() != embd.shape[1]:
+                    logger.warning(f"DFlash: mask embedding width {mask.numel()} != token_embd width {embd.shape[1]}, skipping")
+                elif not (0 <= int(mask_token_id) < embd.shape[0]):
+                    logger.warning(f"DFlash: mask_token_id {mask_token_id} out of range for vocab size {embd.shape[0]}, skipping")
+                else:
+                    embd[int(mask_token_id)] = mask.to(embd.dtype)
+                    logger.info(f"DFlash: synthesized token_embd with mask row {mask_token_id} from {mask_emb_file.name}")
+                    yield "model.embed_tokens.weight", embd
+        yield from super().generate_extra_tensors()
+
+    def _load_target_embed_tokens(self) -> Tensor | None:
+        """Load the target model's token embedding table from target_model_dir."""
+        from safetensors import safe_open
+        target_dir = self.target_model_dir
+        index_file = target_dir / "model.safetensors.index.json"
+        embed_name: str | None = None
+        shard: str | None = None
+        if index_file.is_file():
+            with open(index_file, "r", encoding="utf-8") as f:
+                weight_map = json.load(f).get("weight_map", {})
+            for cand in ("model.embed_tokens.weight", "model.language_model.embed_tokens.weight", "embed_tokens.weight"):
+                if cand in weight_map:
+                    embed_name, shard = cand, weight_map[cand]
+                    break
+            if embed_name is None:
+                for k in sorted(weight_map):
+                    if k.endswith("embed_tokens.weight"):
+                        embed_name, shard = k, weight_map[k]
+                        break
+        else:
+            single = target_dir / "model.safetensors"
+            if not single.is_file():
+                return None
+            embed_name, shard = "model.embed_tokens.weight", "model.safetensors"
+        if embed_name is None or shard is None:
+            return None
+        with safe_open(target_dir / shard, framework="pt") as f:
+            if embed_name not in f.keys():
+                return None
+            return f.get_tensor(embed_name)
+
    @classmethod
    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
        name, gen = item
@@ -228,12 +228,10 @@ class Qwen3ASRMmprojModel(Qwen3OmniMmprojModel):
@ModelBase.register("Glm4vForConditionalGeneration", "Glm4vMoeForConditionalGeneration", "GlmOcrForConditionalGeneration")
@ModelBase.example("zai-org/GLM-4.1V-9B-Thinking", "zai-org/GLM-4.5V")
 class Glm4VVisionModel(Qwen3VLVisionModel):
-    projector_type = gguf.VisionProjectorType.GLM4V
-
    def set_gguf_parameters(self):
        MmprojModel.set_gguf_parameters(self) # skip Qwen3VLVisionModel parameters
        assert self.hparams_vision is not None
-        self.gguf_writer.add_clip_projector_type(self.projector_type)
+        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.GLM4V)

        hidden_act = str(self.hparams_vision.get("hidden_act", "")).lower()
        if hidden_act == "gelu":
@@ -251,32 +249,6 @@ class Glm4VVisionModel(Qwen3VLVisionModel):
        yield from super().modify_tensors(data_torch, name, bid)


-@ModelBase.register("Glm5NextForConditionalGeneration")
-@ModelBase.example("zai-org/GLM-5.3-Flash")
-class Glm5NextVisionModel(Glm4VVisionModel):
-    # GLM-5.3-Flash vision tower. glm4v layout with per-head qk-norm, no post-conv norm and no learned position embeddings.
-    # Images are placed on a ceil aligned canvas with padding.
-
-    projector_type = gguf.VisionProjectorType.GLM5V
-
-    def set_gguf_parameters(self):
-        super().set_gguf_parameters()
-        assert self.hparams_vision is not None
-        self.gguf_writer.add_vision_spatial_merge_size(int(self.hparams_vision.get("spatial_merge_size", 2)))
-        if (limit := self.hparams_vision.get("swiglu_limit")) is not None:
-            self.gguf_writer.add_vision_swiglu_clamp(float(limit))
-
-        # image token budget from the processor, stored as single-frame pixel counts
-        pc = self.preprocessor_config
-        patch = int(pc.get("patch_size", 14))
-        merge = int(pc.get("merge_size", 2))
-        pixels_per_token = (patch * merge) ** 2
-        if (min_tok := pc.get("min_image_tokens")) is not None:
-            self.gguf_writer.add_vision_min_pixels(int(min_tok) * pixels_per_token)
-        if (max_tok := pc.get("max_image_tokens")) is not None:
-            self.gguf_writer.add_vision_max_pixels(int(max_tok) * pixels_per_token)
-
-
@ModelBase.register("Qwen3VLForConditionalGeneration")
@ModelBase.example("Qwen/Qwen3-VL-4B-Instruct")
 class Qwen3VLTextModel(Qwen3Model):
@@ -25,34 +25,15 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):

    model_arch = gguf.MODEL_ARCH.QWEN4EXP

-    # the MTP head: one full-attention QSA block after the trunk, fed by the trunk's hc-wide residual
-    supports_mtp_export = True
-
-    # MTP tensors the shared Qwen remapper does not know
-    _MTP_EXTRA = {
-        "fc_embedding":           "nextn_fc_embedding",
-        "fc_hidden":              "nextn_fc_hidden",
-        "hyper_connection_mixer": "nextn_hc_head",
-    }
+    # the MTP block is a separate draft head; vLLM drops it too
+    supports_mtp_export = False
+    no_mtp = True

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # only the shard names, so the table itself is never held
        self._ple_shards: dict[int, str] = {}
        self._ple_row_dim: int | None = None
-        self._mtp_fc: dict[str, Tensor] = {}
-
-    @classmethod
-    def filter_tensors(cls, item):
-        name, gen = item
-        part = name.split(".")[1] if name.startswith("mtp.") else None
-        if part in cls._MTP_EXTRA:
-            if cls.no_mtp:
-                return None
-            assert cls._original_block_count is not None
-            rest = name.split(".", 2)[2]
-            return f"model.layers.{cls._original_block_count}.{cls._MTP_EXTRA[part]}.{rest}", gen
-        return super().filter_tensors(item)

    def _read_hash_constants(self, suffix: str) -> list[int]:
        """Read an int64 PLE constant straight from the checkpoint.
@@ -82,17 +63,14 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
        self.gguf_writer.add_indexer_top_k(hp["indexer_budget"])
        ratio = hp["indexer_compress_ratio"]
        layer_types = hp["layer_types"]
-        # the MTP block is a full-attention QSA layer too
        self.gguf_writer.add_attention_compress_ratios(
            [ratio if layer_types[i] == "full_attention" else 0 for i in range(n_layer)]
-            + [ratio] * (self.block_count - n_layer)
        )

        # ple_layer_ids is 1-based in the HF config; empty means no n-gram table,
        # so emit no PLE keys rather than optional ones
-        # the MTP head never reads PLE, so an MTP-only file carries none of it
        ple_layers = [i - 1 for i in hp["ple_layer_ids"]]
-        if not ple_layers or self.mtp_only:
+        if not ple_layers:
            return
        self.gguf_writer.add_ple_layers(ple_layers)
        self.gguf_writer.add_ple_ngram_size(hp["ngram_size"])
@@ -142,14 +120,6 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
        if ".ngram_embedding.shard_" in name:
            return self._place_ple_shard(data_torch, name)

-        # eh_proj([e ; h_s]) = fc_embedding(e) + fc_hidden(h_s) for every hc stream s
-        if name.endswith((".nextn_fc_embedding.weight", ".nextn_fc_hidden.weight")):
-            self._mtp_fc[name.rsplit(".", 2)[1]] = data_torch
-            if len(self._mtp_fc) < 2:
-                return []
-            eh = torch.cat([self._mtp_fc.pop("nextn_fc_embedding"), self._mtp_fc.pop("nextn_fc_hidden")], dim=1)
-            return [(self.format_tensor_name(gguf.MODEL_TENSOR.NEXTN_EH_PROJ, bid, ".weight"), eh)]
-
        # one projection feeds indexer q and k; split it, as minimax-m3 does
        if ".indexer.index_qk_proj.weight" in name:
            n_q = self.hparams["indexer_n_heads"] * self.hparams["indexer_head_dim"]
@@ -212,8 +182,6 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):

    def prepare_tensors(self):
        super().prepare_tensors()
-        if self._mtp_fc:
-            raise ValueError(f"MTP projection missing its other half: {sorted(self._mtp_fc)}")
        n_parts = self.hparams.get("split_ngram_parts", 0)
        if self._ple_shards and len(self._ple_shards) != n_parts:
            raise ValueError(
@@ -509,7 +509,6 @@ The following templates have active tests in `tests/test-chat.cpp`:
 | Kimi-K2 / Kimi-K2-Instruct | JSON_NATIVE | JSON tools with special markers |
 | Llama 3.1/3.2/3.3 | JSON_NATIVE | Standard Llama tool format |
 | OpenAI GPT-OSS | Specialized | Channel-based (dedicated handler) |
-| LLM-jp-4.1 | Specialized | GPT-OSS dialect (dedicated handler) |
 | Apriel 1.5 | JSON_NATIVE | `<tool_calls>` wrapper with JSON array |
 | Apriel 1.6 Thinker | Reasoning | Implicit reasoning start |
 | Mistral Small 3.2 | JSON_NATIVE | `[TOOL_CALLS]func[ARGS]{...}` with call ID |
@@ -1,140 +0,0 @@
-# AMD AOCL-BLAS
-
-> [!NOTE]
-> The [ZenDNN backend](ZenDNN.md) is the recommended path for inference on AMD CPUs. Refer to its documentation for the currently supported operations and data types. This page covers AOCL-BLAS as a vendor option for the generic `GGML_BLAS` backend.
-
-AOCL-BLAS is AMD's BLAS library, optimized for AMD EPYC and Ryzen CPUs.
-llama.cpp can link against it through the existing BLAS backend (`GGML_BLAS`).
-
-The BLAS backend can use AOCL-BLAS for eligible large prompt GEMMs and generally does not participate in token generation.
-F32 weights are passed to `cblas_sgemm` directly. Other types are converted to F32 first, so a quantized model can be slower than the native CPU kernels.
-See [BLAS Build](../build.md#blas-build).
-
-AOCL download / install: https://www.amd.com/en/developer/aocl.html
-
-Use the Quick start for a short command list. The later sections cover install layout, the single-threaded tree, threading, and checks.
-
-### Quick start (MT, 64 threads)
-
-If AOCL is not installed yet, follow [Prepare](#prepare) first. Adjust the `amd-libs.cfg` path to your install. CMake flags are set once. Source `amd-libs.cfg` again in every new shell before you launch, or the loader will not find AOCL-BLAS.
-
-```bash
-source /opt/aocl/<version>/aocc/MT/amd-libs.cfg
-
-cmake -B build \
-  -DGGML_BLAS=ON \
-  -DGGML_BLAS_VENDOR=AOCL_mt \
-  -DBLAS_INCLUDE_DIRS="${AOCL_ROOT}/include" \
-  -DGGML_NATIVE=ON
-cmake --build build --config Release
-
-source /opt/aocl/<version>/aocc/MT/amd-libs.cfg
-./build/bin/llama-cli -m model.gguf -t 64
-```
-
-`-t 64` is the thread count to pass on each launch (`--threads 64` is the same flag). Details, the single-threaded tree, and NUMA binding are below.
-
-CMake vendor names `AOCL` and `AOCL_mt` are recognized by [FindBLAS](https://cmake.org/cmake/help/latest/module/FindBLAS.html#blas-lapack-vendors) (CMake 3.27+).
-When either vendor is selected, llama.cpp enables the BLIS code path (`GGML_BLAS_USE_BLIS`): it includes `blis.h` and calls `bli_thread_set_num_threads()` before GEMM. The same vendors also set `GGML_BLAS_USE_AOCL`, which only changes the device description to `AOCL-BLAS`. Upstream BLIS (`FLAME`) sets `GGML_BLAS_USE_BLIS` alone, so its description stays `BLIS`.
-
-### Prepare
-
-1. Install AOCL from AMD (package or tarball). Current releases ship **ST** (single-threaded) and **MT** (multi-threaded) libraries in separate folders. The default `AOCL_ROOT` is the **MT** tree. A typical layout (replace `<version>` and `aocc` with your install):
-
-```
-<aocl-prefix>/<version>/aocc/MT/
-<aocl-prefix>/<version>/aocc/ST/
-```
-
-2. Source `amd-libs.cfg` from the tree you want. Current AOCL versions use this file (not a separate `aocl-env.sh`). Adjust the prefix, version, and compiler (`aocc` vs `gcc`) to match your install:
-
-```bash
-# Multi-threaded (default AOCL_ROOT):
-source /opt/aocl/<version>/aocc/MT/amd-libs.cfg
-
-# Single-threaded:
-# source /opt/aocl/<version>/aocc/ST/amd-libs.cfg
-```
-
-This sets library and include paths so the linker can find AOCL-BLAS. Skipping it is a common cause of BLAS not found / unresolved symbol errors.
-
-Optional, if your install provides an environment module:
-
-```bash
-cd /opt/aocl/<version>/aocc/MT
-module load ./aocl-linux-aocc-<version>_module
-# module unload ./aocl-linux-aocc-<version>_module
-```
-
-3. Prefer the **MT** libraries for llama.cpp. Use `-DGGML_BLAS_VENDOR=AOCL_mt` after sourcing the MT `amd-libs.cfg`. Use `-DGGML_BLAS_VENDOR=AOCL` if you sourced the ST tree instead.
-
-### llama.cpp compilation
-
-Requires **CMake 3.27 or newer** for `-DGGML_BLAS_VENDOR=AOCL` / `AOCL_mt`.
-
-FindBLAS does not detect AOCL headers via pkg-config. After sourcing `amd-libs.cfg`, `AOCL_ROOT` is set and `$AOCL_ROOT/include` is a symlink to the active integer ABI (`include_LP64` by default). Pass that path to CMake:
-
-```bash
-source /opt/aocl/<version>/aocc/MT/amd-libs.cfg   # adjust path
-
-cmake -B build \
-  -DGGML_BLAS=ON \
-  -DGGML_BLAS_VENDOR=AOCL_mt \
-  -DBLAS_INCLUDE_DIRS="${AOCL_ROOT}/include" \
-  -DGGML_NATIVE=ON
-
-cmake --build build --config Release
-```
-
-#### CMake older than 3.27
-
-`AOCL` / `AOCL_mt` may be unknown to FindBLAS. After sourcing the AOCL env, you can try:
-
-```bash
-cmake -B build \
-  -DGGML_BLAS=ON \
-  -DGGML_BLAS_VENDOR=Generic \
-  -DBLAS_LIBRARIES="-lblis -lm" \
-  -DBLAS_INCLUDE_DIRS="${AOCL_ROOT}/include" \
-  -DGGML_NATIVE=ON
-```
-
-Library names differ between AOCL packages (`blis`, `blis-mt`, etc.). Pass whatever your install provides.
-`GGML_BLAS_VENDOR=Generic` does not enable the BLIS header and thread path (`GGML_BLAS_USE_BLIS`). Upgrade CMake so `AOCL` or `AOCL_mt` is recognized.
-
-### llama.cpp execution
-
-`--threads` / `--threads-batch` are the thread budget for **every** backend that implements `set_n_threads` (CPU and BLAS).
-
-On each large BLAS `MUL_MAT`, the backend calls `bli_thread_set_num_threads()` with that **same** value, so BLIS GEMM may use up to `--threads-batch` threads during prompt processing. Other ops stay on the CPU backend with the same limit. Token generation usually does not use BLAS.
-
-`BLIS_NUM_THREADS` is **not** a reliable way to cap BLIS here: the per-GEMM `bli_thread_set_num_threads()` call overrides it. To use fewer cores, lower `--threads` and/or `--threads-batch`.
-
-Thread scaling depends on the CPU, NUMA layout, model, and batch size. Benchmark the thread counts used for deployment. Nested OpenMP (ggml type conversion, then BLIS GEMM, both using OpenMP) can still oversubscribe even though those two steps are sequential.
-
-On a multi-socket machine, bind the process to one NUMA node. To skip SMT, bind to that node's physical cores only (check `lscpu -e`; on many AMD layouts the first range is the physical cores and a higher range is the sibling threads):
-
-```bash
-numactl --physcpubind=0-127 --membind=0 ./build/bin/llama-cli -m model.gguf
-```
-
-Keep the sourced AOCL env (or `LD_LIBRARY_PATH`) set when running binaries, or dynamic linking to AOCL libs will fail. Source the same `amd-libs.cfg` you used at build time.
-
-### Verify
-
- Configure output should show BLAS found, with libraries under the AOCL MT tree and includes at `$AOCL_ROOT/include`.
- `ldd` on `llama-bench` should list that same AOCL-BLAS library; `blis-mt` or `blis`.
- `--list-devices` prints the description `AOCL-BLAS` (`BLAS: AOCL-BLAS`). Upstream BLIS (`FLAME`) still prints `BLIS`.
- `llama-bench` reports the backend as `BLAS` for this build and `CPU` for a build with `-DGGML_BLAS=OFF`.
- `test-backend-ops -b BLAS` checks that BLAS GEMMs match the CPU reference.
-
-### Notes
-
- Optional `-march=znver3` / `znver4` / `znver5` (or similar) can be passed via `CMAKE_C_FLAGS` / `CMAKE_CXX_FLAGS` for a known CPU, but `-DGGML_NATIVE=ON` is usually enough and is safer across Ryzen / EPYC generations.
- For building AOCL-BLAS (AMD's BLIS fork) from source instead of AOCL packages, see https://github.com/amd/blis
-
-### Reference
-
-1. https://www.amd.com/en/developer/aocl.html
-2. https://cmake.org/cmake/help/latest/module/FindBLAS.html#blas-lapack-vendors
-3. https://github.com/amd/blis
@@ -52,10 +52,6 @@ The packages for FP32 and FP16 would have different accuracy and performance on

 ## News

- 2026.09
-  - Update the CI build environment for oneAPI 2026.1 (unified oneAPI Toolkit). oneDNN is removed from the Deep Learning Essentials package in 2026.0, so the CI now uses the oneAPI Toolkit installer which still includes oneDNN.
-  - oneAPI 2026.1 improves the SYCL build performance: measured with the same code on Arc B570, prompt processing 1331 vs 434 t/s (3.1x) vs the 2025.3-based release build.
-
 - 2026.04-05
  - Optimize mul_mat by reorder feature for data type: Q4_K, Q5_K, Q6_K, Q8_0.
  - Fused MoE.
@@ -261,7 +257,7 @@ Platform #0: Intel(R) OpenCL HD Graphics
 `-- Device #0: Intel(R) Iris(R) Xe Graphics [0x9a49]
 ```

-2. **Install Intel® oneAPI Toolkit**
+2. **Install Intel® oneAPI Base toolkit**

 SYCL backend depends on:
  - Intel® oneAPI DPC++/C++ compiler/running-time.
@@ -271,11 +267,11 @@ SYCL backend depends on:

 - **For Intel GPU**

-With the 2026.0 release, the Intel® oneAPI Base toolkit and the HPC toolkit are combined into the **Intel® oneAPI Toolkit**, and **oneDNN is removed from the Intel® Deep Learning Essentials** package (oneDNN is distributed separately since then). The **Intel® oneAPI Toolkit** includes oneDNN until 2027.0.
+All above are included in both **Intel® oneAPI Base toolkit** and **Intel® Deep Learning Essentials** packages.

-It's recommended to install the **Intel® oneAPI Toolkit**.
+It's recommended to install **Intel® Deep Learning Essentials** which only provides the necessary libraries with less size.

-The **Intel® oneAPI Toolkit** can be obtained from the official [Intel® oneAPI Toolkit](https://www.intel.com/content/www/us/en/developer/tools/oneapi/base-toolkit-download.html) page.
+The **Intel® oneAPI Base toolkit** and **Intel® Deep Learning Essentials** can be obtained from the official [Intel® oneAPI Base Toolkit](https://www.intel.com/content/www/us/en/developer/tools/oneapi/base-toolkit.html) page.

 Please follow the instructions for downloading and installing the Toolkit for Linux, and preferably keep the default installation values unchanged, notably the installation path *(`/opt/intel/oneapi` by default)*.

@@ -285,7 +281,6 @@ Upon a successful installation, SYCL is enabled for the available Intel devices,

 |Verified release|
 |-|
-|2026.1 |
 |2025.3.3 |
 |2025.2.1|
 |2025.1|
@@ -116,20 +116,6 @@ This provides BLAS acceleration using only the CPU. Make sure to have OpenBLAS i

 Check [BLIS.md](./backend/BLIS.md) for more information.

-### AMD AOCL-BLAS
-
-For AMD CPU inference, the [ZenDNN backend](#zendnn) is recommended. AOCL-BLAS is also available as a vendor option for the generic `GGML_BLAS` backend.
-
-Source `amd-libs.cfg` from your AOCL install (MT tree by default), then build (CMake 3.27+ recommended for the `AOCL` / `AOCL_mt` vendors):
-
-```bash
-source /opt/aocl/<version>/aocc/MT/amd-libs.cfg   # adjust path; ST tree uses .../ST/amd-libs.cfg
-cmake -B build -DGGML_BLAS=ON -DGGML_BLAS_VENDOR=AOCL_mt -DBLAS_INCLUDE_DIRS="${AOCL_ROOT}/include" -DGGML_NATIVE=ON
-cmake --build build --config Release
-```
-
-Full steps, threading notes, and a fallback for older CMake: [AOCL.md](./backend/AOCL.md).
-
 ### Intel oneMKL

 Building through oneAPI compilers will make avx_vnni instruction set available for intel processors that do not support avx512 and avx512_vnni. Please note that this build config **does not support Intel GPU**. For Intel GPU support, please refer to [llama.cpp for SYCL](./backend/SYCL.md).
@@ -195,8 +181,6 @@ cmake -B build -DGGML_CUDA=ON
 cmake --build build --config Release
 ```

-To use a specific CCCL version instead of the one bundled with the installed CUDA Toolkit, add `-DGGML_CUDA_CCCL_VERSION=vMAJOR.MINOR.PATCH`. CUB DeviceTopK requires CCCL 3.4.3 or newer; older versions use the sort fallback.
-
 Note that this also builds the CPU backend by default. On Windows on ARM, MSVC's
 support for the ARM NEON intrinsics used by the CPU backend may be incomplete, so
 a CUDA build produced entirely with MSVC might have a slower CPU backend. If CPU
@@ -16,7 +16,6 @@ Function calling is supported for all models (see https://github.com/ggml-org/ll
  - Firefunction v2
  - Command R7B
  - DeepSeek R1 (WIP / seems reluctant to call any tools?)
-  - GPT-OSS (Harmony), LLM-jp-4.1 (Harmony dialect)

 - Generic tool call is supported when the template isn't recognized by native format handlers (you'll see `Chat format: Generic` in the logs).
  - Use `--chat-template-file` to override the template when appropriate (see examples below)
@@ -17,17 +17,17 @@ Legend:
 |                              ABS | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                              ACC | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              ADD | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             ADD1 | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                             ADD1 | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                           ADD_ID | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                           ARANGE | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                           ARGMAX | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          ARGSORT | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             CEIL | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                            CLAMP | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                        COL2IM_1D | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                        COL2IM_1D | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                           CONCAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                             CONT | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
-|                          CONV_2D | ❌ | ❌ | 🟡 | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ |
+|                          CONV_2D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ |
 |                       CONV_2D_DW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          CONV_3D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                CONV_TRANSPOSE_1D | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -41,9 +41,9 @@ Legend:
 |                             DIAG | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                    DIAG_MASK_INF | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              DIV | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                     DSV4_HC_COMB | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                     DSV4_HC_POST | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                      DSV4_HC_PRE | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                     DSV4_HC_COMB | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                     DSV4_HC_POST | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                      DSV4_HC_PRE | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                              DUP | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              ELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                              EXP | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
@@ -68,17 +68,17 @@ Legend:
 |                        IM2COL_3D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                          L2_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                       LEAKY_RELU | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | 🟡 | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                LIGHTNING_INDEXER | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                LIGHTNING_INDEXER | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                              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 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                  MUL_MAT_ID_W4A4 | ❌ | ❌ | ❌ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                  MUL_MAT_ID_W4A8 | ❌ | ❌ | ❌ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                     MUL_MAT_W4A4 | ❌ | ❌ | ❌ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                     MUL_MAT_W4A8 | ❌ | ❌ | ❌ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                              NEG | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                   OPT_STEP_ADAMW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -121,7 +121,7 @@ Legend:
 |                              SUM | ❌ | 🟡 | ✅ | 🟡 | ❌ | 🟡 | 🟡 | ❌ | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
 |                         SUM_ROWS | ❌ | ✅ | ✅ | 🟡 | ❌ | 🟡 | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ |
 |                           SWIGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                     SWIGLU_CLAMP | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                     SWIGLU_CLAMP | ❌ | ❌ | ❌ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                       SWIGLU_OAI | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             TANH | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |               TIMESTEP_EMBEDDING | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -117,7 +117,7 @@ int main(int argc, char ** argv) {

    // create a llama_batch
    // we use this object to submit token data for decoding
-    common_batch batch(ctx);
+    llama_batch batch = llama_batch_init(std::max(tokens_list.size(), (size_t) n_parallel), 0, n_parallel);

    std::vector<llama_seq_id> seq_ids(n_parallel, 0);
    for (int32_t i = 0; i < n_parallel; ++i) {
@@ -126,12 +126,12 @@ int main(int argc, char ** argv) {

    // evaluate the initial prompt
    for (size_t i = 0; i < tokens_list.size(); ++i) {
-        batch.add(tokens_list[i], i, seq_ids, false);
+        common_batch_add(batch, tokens_list[i], i, seq_ids, false);
    }
-    GGML_ASSERT(batch.size() == (int) tokens_list.size());
+    GGML_ASSERT(batch.n_tokens == (int) tokens_list.size());

    if (llama_model_has_encoder(model)) {
-        if (llama_process(ctx, LLAMA_PROCESS_TYPE_ENCODE, batch.get())) {
+        if (llama_encode(ctx, batch)) {
            LOG_ERR("%s : failed to eval\n", __func__);
            return 1;
        }
@@ -141,14 +141,14 @@ int main(int argc, char ** argv) {
            decoder_start_token_id = llama_vocab_bos(vocab);
        }

-        batch.clear();
-        batch.add(decoder_start_token_id, 0, seq_ids, false);
+        common_batch_clear(batch);
+        common_batch_add(batch, decoder_start_token_id, 0, seq_ids, false);
    }

    // llama_decode will output logits only for the last token of the prompt
-    batch.set_output(batch.size() - 1, true);
+    batch.logits[batch.n_tokens - 1] = true;

-    if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()) != 0) {
+    if (llama_decode(ctx, batch) != 0) {
        LOG_ERR("%s: llama_decode() failed\n", __func__);
        return 1;
    }
@@ -170,16 +170,16 @@ int main(int argc, char ** argv) {

    // remember the batch index of the last token for each parallel sequence
    // we need this to determine which logits to sample from
-    std::vector<int32_t> i_batch(n_parallel, batch.size() - 1);
+    std::vector<int32_t> i_batch(n_parallel, batch.n_tokens - 1);

-    int n_cur    = batch.size();
+    int n_cur    = batch.n_tokens;
    int n_decode = 0;

    const auto t_main_start = ggml_time_us();

    while (n_cur <= n_predict) {
        // prepare the next batch
-        batch.clear();
+        common_batch_clear(batch);

        // sample the next token for each parallel sequence / stream
        for (int32_t i = 0; i < n_parallel; ++i) {
@@ -208,23 +208,23 @@ int main(int argc, char ** argv) {

            streams[i] += common_token_to_piece(ctx, new_token_id);

-            i_batch[i] = batch.size();
+            i_batch[i] = batch.n_tokens;

            // push this new token for next evaluation
-            batch.add(new_token_id, n_cur, i, true);
+            common_batch_add(batch, new_token_id, n_cur, { i }, true);

            n_decode += 1;
        }

        // all streams are finished
-        if (batch.size() == 0) {
+        if (batch.n_tokens == 0) {
            break;
        }

        n_cur += 1;

        // evaluate the current batch with the transformer model
-        if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
+        if (llama_decode(ctx, batch)) {
            LOG_ERR("%s : failed to eval, return code %d\n", __func__, 1);
            return 1;
        }
@@ -249,6 +249,7 @@ int main(int argc, char ** argv) {

    fprintf(stderr, "\n");

+    llama_batch_free(batch);

    for (auto & sampler_config : sampler_configs) {
        llama_sampler_free(sampler_config.sampler);
@@ -194,8 +194,7 @@ static bool run(llama_context * ctx, const common_params & params) {
        return false;
    }

-    common_batch batch = common_batch_get_one(ctx, tokens);
-    if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
+    if (llama_decode(ctx, llama_batch_get_one(tokens.data(), tokens.size()))) {
        LOG_ERR("%s : failed to eval\n", __func__);
        return false;
    }
@@ -1,7 +1,5 @@
 #include "diffusion.h"

-#include "common.h"
-
 #include "log.h"

 #include <algorithm>
@@ -146,7 +144,8 @@ void diffusion_generate(llama_context *          ctx,

    struct llama_sampler * dist_sampler = llama_sampler_init_dist(params.seed);

-    common_batch batch(ctx);
+    llama_batch batch = llama_batch_init(params.max_length, 0, 1);
+    batch.n_tokens    = params.max_length;

    // Pre-allocate buffers for CFG if needed
    int32_t                  logits_size = n_vocab * params.max_length;
@@ -203,15 +202,18 @@ void diffusion_generate(llama_context *          ctx,
            }

            // Setup batch
-            batch.clear();
            for (int32_t i = 0; i < params.max_length; i++) {
-                batch.add(output_tokens[i], i, 0, true);
+                batch.token[i]     = output_tokens[i];
+                batch.pos[i]       = i;
+                batch.n_seq_id[i]  = 1;
+                batch.seq_id[i][0] = 0;
+                batch.logits[i]    = 1;
            }

            float * logits = nullptr;

            if (params.cfg_scale > 0.0f) {
-                int ret = llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get());
+                int ret = llama_decode(ctx, batch);
                if (ret != 0) {
                    LOG_ERR("Failed to generate conditional");
                    break;
@@ -225,11 +227,10 @@ void diffusion_generate(llama_context *          ctx,
                    un_x_buffer[i] = params.mask_token_id;
                }

-                batch.clear();
                for (int32_t i = 0; i < params.max_length; i++) {
-                    batch.add(un_x_buffer[i], i, 0, true);
+                    batch.token[i] = un_x_buffer[i];
                }
-                ret = llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get());
+                ret = llama_decode(ctx, batch);
                if (ret != 0) {
                    LOG_ERR("Failed to generate unconditional");
                    break;
@@ -243,7 +244,7 @@ void diffusion_generate(llama_context *          ctx,
                }
                logits = cond_logits_buffer.data();
            } else {
-                int ret = llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get());
+                int ret = llama_decode(ctx, batch);
                if (ret != 0) {
                    LOG_ERR("%s: failed to decode at step %d, ret = %d\n", __func__, global_step, ret);
                    break;
@@ -399,6 +400,7 @@ void diffusion_generate(llama_context *          ctx,
            total_time / 1000.0 / params.steps,
            total_sampling_time / 1000.0 / params.steps);

+    llama_batch_free(batch);
    llama_sampler_free(sampler);
    llama_sampler_free(dist_sampler);

@@ -27,27 +27,27 @@ static std::vector<std::string> split_lines(const std::string & s, const std::st
    return lines;
 }

-static void batch_add_seq(common_batch & batch, const std::vector<int32_t> & tokens, llama_seq_id seq_id) {
+static void batch_add_seq(llama_batch & batch, const std::vector<int32_t> & tokens, llama_seq_id seq_id) {
    size_t n_tokens = tokens.size();
    for (size_t i = 0; i < n_tokens; i++) {
-        batch.add(tokens[i], i, seq_id, true);
+        common_batch_add(batch, tokens[i], i, { seq_id }, true);
    }
 }

-static void batch_decode(llama_context * ctx, common_batch & batch, float * output, int n_seq, int n_embd_out, int embd_norm) {
+static void batch_decode(llama_context * ctx, llama_batch & batch, float * output, int n_seq, int n_embd_out, int embd_norm) {
    const enum llama_pooling_type pooling_type = llama_pooling_type(ctx);

    // clear previous kv_cache values (irrelevant for embeddings)
    llama_memory_clear(llama_get_memory(ctx), true);

    // run model
-    LOG_INF("%s: n_tokens = %d, n_seq = %d\n", __func__, batch.size(), n_seq);
-    if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()) < 0) {
+    LOG_INF("%s: n_tokens = %d, n_seq = %d\n", __func__, batch.n_tokens, n_seq);
+    if (llama_decode(ctx, batch) < 0) {
        LOG_ERR("%s : failed to process\n", __func__);
    }

-    for (int i = 0; i < batch.size(); i++) {
-        if (!batch.tokens[i].output) {
+    for (int i = 0; i < batch.n_tokens; i++) {
+        if (!batch.logits[i]) {
            continue;
        }

@@ -61,8 +61,8 @@ static void batch_decode(llama_context * ctx, common_batch & batch, float * outp
            GGML_ASSERT(embd != NULL && "failed to get token embeddings");
        } else {
            // try to get sequence embeddings - supported only when pooling_type is not NONE
-            embd = llama_get_embeddings_seq(ctx, batch.tokens[i].seq_id);
-            embd_pos = batch.tokens[i].seq_id;
+            embd = llama_get_embeddings_seq(ctx, batch.seq_id[i][0]);
+            embd_pos = batch.seq_id[i][0];
            GGML_ASSERT(embd != NULL && "failed to get sequence embeddings");
        }

@@ -242,7 +242,7 @@ int main(int argc, char ** argv) {

    // initialize batch
    const int n_prompts = prompts.size();
-    common_batch batch(ctx);
+    struct llama_batch batch = llama_batch_init(n_batch, 0, 1);

    // count number of embeddings
    int n_embd_count = 0;
@@ -269,12 +269,12 @@ int main(int argc, char ** argv) {
        const uint64_t n_toks = inp.size();

        // encode if at capacity
-        if (batch.size() + n_toks > n_batch || s >= n_seq_max) {
+        if (batch.n_tokens + n_toks > n_batch || s >= n_seq_max) {
            float * out = emb + e * n_embd_out;
            batch_decode(ctx, batch, out, s, n_embd_out, params.embd_normalize);
-            e += pooling_type == LLAMA_POOLING_TYPE_NONE ? batch.size() : s;
+            e += pooling_type == LLAMA_POOLING_TYPE_NONE ? batch.n_tokens : s;
            s = 0;
-            batch.clear();
+            common_batch_clear(batch);
        }

        // add to batch
@@ -407,6 +407,7 @@ int main(int argc, char ** argv) {
    llama_perf_context_print(ctx);

    // clean up
+    llama_batch_free(batch);
    llama_backend_free();

    return 0;
@@ -26,8 +26,7 @@ static bool run(llama_context * ctx, const common_params & params) {
        LOG_INF("  %d\n", tokens[i]);
    }

-    common_batch batch = common_batch_get_one(ctx, tokens);
-    if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
+    if (llama_decode(ctx, llama_batch_get_one(tokens.data(), tokens.size()))) {
        LOG_ERR("%s : failed to eval\n", __func__);
        return false;
    }
@@ -57,13 +57,12 @@ int main(int argc, char ** argv) {
        return 1;
    }

+    llama_batch batch = llama_batch_get_one(prompt_tokens.data(), prompt_tokens.size());
+
    const int n_iters = 3;

    // warm-up
-    {
-        common_batch batch = common_batch_get_one(ctx, prompt_tokens);
-        llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get());
-    }
+    llama_decode(ctx, batch);
    llama_memory_clear(llama_get_memory(ctx), true);
    llama_synchronize(ctx);

@@ -72,16 +71,13 @@ int main(int argc, char ** argv) {
        double t_sum2_us = 0.0;

        for (int i = 0; i < n_iters; i++) {
-            // positions continue from the memory
-            common_batch batch = common_batch_get_one(ctx, prompt_tokens);
-
            // this pause is important - it simulates "idle GPU"
            std::this_thread::sleep_for(std::chrono::milliseconds(t_pause_ms));

            const int64_t t_start_us = llama_time_us();

            // this should take constant time
-            llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get());
+            llama_decode(ctx, batch);
            llama_synchronize(ctx);

            const int64_t t_end_us = llama_time_us();
@@ -35,7 +35,7 @@ constexpr float DEFAULT_SAMPLER_TEMP    = 0.3f;

 static llama_model                      * g_model;
 static llama_context                    * g_context;
-static common_batch                       g_batch;
+static llama_batch                        g_batch;
 static common_chat_templates_ptr          g_chat_templates;
 static common_sampler                   * g_sampler;

@@ -116,7 +116,7 @@ Java_com_arm_aichat_internal_InferenceEngineImpl_prepare(JNIEnv * /*env*/, jobje
    auto *context = init_context(g_model);
    if (!context) { return 1; }
    g_context = context;
-    g_batch = common_batch(context);
+    g_batch = llama_batch_init(BATCH_SIZE, 0, 1);
    g_chat_templates = common_chat_templates_init(g_model, "");
    g_sampler = new_sampler(DEFAULT_SAMPLER_TEMP);
    return 0;
@@ -164,18 +164,18 @@ Java_com_arm_aichat_internal_InferenceEngineImpl_benchModel(JNIEnv *env, jobject
    for (nri = 0; nri < nr; nri++) {
        LOGi("Benchmark prompt processing (pp = %d)", pp);

-        common_batch batch(context);
+        common_batch_clear(g_batch);

        const int n_tokens = pp;
        for (i = 0; i < n_tokens; i++) {
-            batch.add(0, i, 0, false);
+            common_batch_add(g_batch, 0, i, {0}, false);
        }

-        batch.set_output(batch.size() - 1, true);
+        g_batch.logits[g_batch.n_tokens - 1] = true;
        llama_memory_clear(llama_get_memory(context), false);

        const auto t_pp_start = ggml_time_us();
-        if (llama_process(context, LLAMA_PROCESS_TYPE_DECODE, batch.get()) != 0) {
+        if (llama_decode(context, g_batch) != 0) {
            LOGe("llama_decode() failed during prompt processing");
        }
        const auto t_pp_end = ggml_time_us();
@@ -187,12 +187,12 @@ Java_com_arm_aichat_internal_InferenceEngineImpl_benchModel(JNIEnv *env, jobject
        llama_memory_clear(llama_get_memory(context), false);
        const auto t_tg_start = ggml_time_us();
        for (i = 0; i < tg; i++) {
-            batch.clear();
+            common_batch_clear(g_batch);
            for (j = 0; j < pl; j++) {
-                batch.add(0, i, j, true);
+                common_batch_add(g_batch, 0, i, {j}, true);
            }

-            if (llama_process(context, LLAMA_PROCESS_TYPE_DECODE, batch.get()) != 0) {
+            if (llama_decode(context, g_batch) != 0) {
                LOGe("llama_decode() failed during text generation");
            }
        }
@@ -315,7 +315,7 @@ static void reset_short_term_states() {

 static int decode_tokens_in_batches(
        llama_context *context,
-        common_batch &batch,
+        llama_batch &batch,
        const llama_tokens &tokens,
        const llama_pos start_pos,
        const bool compute_last_logit = false) {
@@ -323,7 +323,7 @@ static int decode_tokens_in_batches(
    LOGd("%s: Decode %d tokens starting at position %d", __func__, (int) tokens.size(), start_pos);
    for (int i = 0; i < (int) tokens.size(); i += BATCH_SIZE) {
        const int cur_batch_size = std::min((int) tokens.size() - i, BATCH_SIZE);
-        batch.clear();
+        common_batch_clear(batch);
        LOGv("%s: Preparing a batch size of %d starting at: %d", __func__, cur_batch_size, i);

        // Shift context if current batch cannot fit into the context
@@ -337,11 +337,11 @@ static int decode_tokens_in_batches(
            const llama_token token_id = tokens[i + j];
            const llama_pos position = start_pos + i + j;
            const bool want_logit = compute_last_logit && (i + j == tokens.size() - 1);
-            batch.add(token_id, position, 0, want_logit);
+            common_batch_add(batch, token_id, position, {0}, want_logit);
        }

        // Decode this batch
-        const int decode_result = llama_process(context, LLAMA_PROCESS_TYPE_DECODE, batch.get());
+        const int decode_result = llama_decode(context, batch);
        if (decode_result) {
            LOGe("%s: llama_decode failed w/ %d", __func__, decode_result);
            return 1;
@@ -506,9 +506,9 @@ Java_com_arm_aichat_internal_InferenceEngineImpl_generateNextToken(
    common_sampler_accept(g_sampler, new_token_id, true);

    // Populate the batch with new token, then decode
-    g_batch.clear();
-    g_batch.add(new_token_id, current_position, 0, true);
-    if (llama_process(g_context, LLAMA_PROCESS_TYPE_DECODE, g_batch.get()) != 0) {
+    common_batch_clear(g_batch);
+    common_batch_add(g_batch, new_token_id, current_position, {0}, true);
+    if (llama_decode(g_context, g_batch) != 0) {
        LOGe("%s: llama_decode() failed for generated token", __func__);
        return nullptr;
    }
@@ -553,7 +553,7 @@ Java_com_arm_aichat_internal_InferenceEngineImpl_unload(JNIEnv * /*unused*/, job
    // Free up resources
    common_sampler_free(g_sampler);
    g_chat_templates.reset();
-    g_batch = common_batch();
+    llama_batch_free(g_batch);
    llama_free(g_context);
    llama_model_free(g_model);
 }
@@ -101,13 +101,8 @@ int main(int argc, char ** argv) {
    const auto t_enc_start = ggml_time_us();

    // eval the prompt
-    {
-        common_batch batch = common_batch_get_one(ctx, inp.data(), n_input - 1);
-        llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get());
-
-        batch = common_batch_get_one(ctx, &inp.back(), 1);
-        llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get());
-    }
+    llama_decode(ctx, llama_batch_get_one( inp.data(), n_input - 1));
+    llama_decode(ctx, llama_batch_get_one(&inp.back(),           1));

    for (int s = 1; s < W + G + 1; ++s) {
        llama_memory_seq_cp(mem, 0, s, -1, -1);
@@ -129,7 +124,7 @@ int main(int argc, char ** argv) {
    // seq_id == 0           : the current input token
    // seq_id [1, W]         : tokens from the past N - 1 Jacobi iterations
    // seq_id [W + 1, W + G] : verification n-grams
-    common_batch batch(ctx);
+    llama_batch batch = llama_batch_init(llama_n_ctx(ctx), 0, W + G + 1);

    // target model sampling context
    struct common_sampler * smpl = common_sampler_init(model, params.sampling);
@@ -209,10 +204,10 @@ int main(int argc, char ** argv) {
        //                                                      V  V  V  V  V  V
        //                                                             id
        {
-            batch.clear();
+            common_batch_clear(batch);

            // current token - first token of the first level
-            batch.add(id, n_past, seq_id_all, true);
+            common_batch_add(batch, id, n_past, seq_id_all, true);

            // verification n-grams - queue this before the lookahead tokens for less KV cache fragmentation
            {
@@ -235,9 +230,9 @@ int main(int argc, char ** argv) {
                        const llama_token t = ngrams_observed.tokens[idx + j];

                        ngrams_cur[g].tokens [j + 1] = t;
-                        ngrams_cur[g].i_batch[j + 1] = batch.size();
+                        ngrams_cur[g].i_batch[j + 1] = batch.n_tokens;

-                        batch.add(t, n_past + j + 1, W + 1 + g, true);
+                        common_batch_add(batch, t, n_past + j + 1, { W + 1 + g }, true);
                    }
                }
            }
@@ -249,18 +244,18 @@ int main(int argc, char ** argv) {
                    seq_id_look[j] = i + j + 1;
                }

-                batch.add(tokens_j[0][i], n_past + i, seq_id_look, false);
+                common_batch_add(batch, tokens_j[0][i], n_past + i, seq_id_look, false);
            }

            // fill the rest of the levels
            for (int j = 1; j < N - 1; j++) {
                for (int i = 0; i < W; i++) {
-                    batch.add(tokens_j[j][i], n_past + j + i, i + 1, j == N - 2);
+                    common_batch_add(batch, tokens_j[j][i], n_past + j + i, { i + 1 }, j == N - 2);
                }
            }
        }

-        if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()) != 0) {
+        if (llama_decode(ctx, batch) != 0) {
            LOG_ERR("\n\n%s: llama_decode failed - increase KV cache size\n", __func__);
            return 1;
        }
@@ -478,6 +473,7 @@ int main(int argc, char ** argv) {

    common_sampler_free(smpl);

+    llama_batch_free(batch);

    llama_backend_free();

@@ -98,13 +98,8 @@ int main(int argc, char ** argv){

    const auto t_enc_start = ggml_time_us();

-    {
-        common_batch batch = common_batch_get_one(ctx, inp.data(), n_input - 1);
-        llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get());
-
-        batch = common_batch_get_one(ctx, &inp.back(), 1);
-        llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get());
-    }
+    llama_decode(ctx, llama_batch_get_one( inp.data(), n_input - 1));
+    llama_decode(ctx, llama_batch_get_one(&inp.back(),           1));

    const auto t_enc_end = ggml_time_us();

@@ -120,7 +115,7 @@ int main(int argc, char ** argv){

    std::vector<llama_token> draft;

-    common_batch batch_tgt(ctx);
+    llama_batch batch_tgt = llama_batch_init(llama_n_ctx(ctx), 0, 1);

    const auto t_dec_start = ggml_time_us();

@@ -197,8 +192,8 @@ int main(int argc, char ** argv){
        // clean the cache of draft tokens that weren't accepted
        llama_memory_seq_rm(llama_get_memory(ctx), 0, n_past, -1);

-        batch_tgt.clear();
-        batch_tgt.add(draft[0], n_past, 0, true);
+        common_batch_clear(batch_tgt);
+        common_batch_add(batch_tgt, draft[0], n_past, { 0 }, true);

        // Draft already contains a single token sampled from the model:
        GGML_ASSERT(draft.size() == 1);
@@ -208,13 +203,13 @@ int main(int argc, char ** argv){
        common_ngram_cache_draft(inp, draft, n_draft, LLAMA_NGRAM_MIN, LLAMA_NGRAM_MAX, ngram_cache_context, ngram_cache_dynamic, ngram_cache_static);

        for (size_t i = 1; i < draft.size(); ++i) {
-            batch_tgt.add(draft[i], n_past + i, 0, true);
+            common_batch_add(batch_tgt, draft[i], n_past + i, { 0 }, true);
        }

        t_draft_us += ggml_time_us() - t_start_draft_us;
        n_drafted += draft.size() - 1;

-        llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch_tgt.get());
+        llama_decode(ctx, batch_tgt);
        ++n_past;

        draft.erase(draft.begin());
@@ -246,6 +241,7 @@ int main(int argc, char ** argv){

    common_sampler_free(smpl);

+    llama_batch_free(batch_tgt);

    llama_backend_free();

@@ -19,8 +19,6 @@ def parse_arguments():
    parser.add_argument("--prompt-file", "-f", help="Optional prompt file", required=False)
    parser.add_argument("--verbose", "-v", action="store_true", help="Enable verbose debug output")
    parser.add_argument("--device", "-d", help="Device to use (cpu, cuda, mps, auto)", default="auto")
-    parser.add_argument("--add-bos", action=argparse.BooleanOptionalAction, default=None,
-                        help="Override BOS token setting (default: use model's own setting)")
    return parser.parse_args()

 def load_model_and_tokenizer(model_path, device="auto"):
@@ -121,9 +119,6 @@ def main():

    model, tokenizer, config = load_model_and_tokenizer(model_path, args.device)

-    if args.add_bos is not None and hasattr(tokenizer, "add_bos_token"):
-        tokenizer.add_bos_token = args.add_bos
-
    if args.verbose:
        enable_torch_debugging(model)

@@ -224,6 +224,8 @@ int main(int argc, char ** argv) {

    LOG_INF("\n\n");

+    const int n_ctx = llama_n_ctx(ctx);
+
    if (sseed >= 0) {
        LOG_INF("%s: initializing all samplers with the same RNG seed: %d (use a negative seed to have different seeds)\n", __func__, sseed);
    } else {
@@ -250,7 +252,7 @@ int main(int argc, char ** argv) {

    // the max batch size is as large as the context to handle cases where we get very long input prompt from multiple
    // users. regardless of the size, the main loop will chunk the batch into a maximum of params.n_batch tokens at a time
-    common_batch batch(ctx);
+    llama_batch batch = llama_batch_init(n_ctx, 0, 1);

    int32_t n_total_prompt = 0;
    int32_t n_total_gen    = 0;
@@ -266,10 +268,10 @@ int main(int argc, char ** argv) {
        LOG_INF("%s: Evaluating the system prompt ...\n", __func__);

        for (int32_t i = 0; i < n_tokens_system; ++i) {
-            batch.add(tokens_system[i], i, 0, false);
+            common_batch_add(batch, tokens_system[i], i, { 0 }, false);
        }

-        if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()) != 0) {
+        if (llama_decode(ctx, batch) != 0) {
            LOG_ERR("%s: llama_decode() failed\n", __func__);
            return 1;
        }
@@ -285,7 +287,7 @@ int main(int argc, char ** argv) {
    LOG_INF("Processing requests ...\n\n");

    while (true) {
-        batch.clear();
+        common_batch_clear(batch);

        // decode any currently ongoing sequences
        for (auto & client : clients) {
@@ -293,14 +295,14 @@ int main(int argc, char ** argv) {
                continue;
            }

-            client.i_batch = batch.size();
+            client.i_batch = batch.n_tokens;

-            batch.add(client.sampled, client.n_past++, client.id + 1, true);
+            common_batch_add(batch, client.sampled, client.n_past++, { client.id + 1 }, true);

            client.n_decoded += 1;
        }

-        if (batch.size() == 0) {
+        if (batch.n_tokens == 0) {
            // all sequences have ended - clear the entire KV cache
            for (int i = 1; i <= n_clients; ++i) {
                llama_memory_seq_rm(mem, i, -1, -1);
@@ -312,7 +314,7 @@ int main(int argc, char ** argv) {
        }

        // insert new sequences for decoding
-        if (cont_batching || batch.size() == 0) {
+        if (cont_batching || batch.n_tokens == 0) {
            for (auto & client : clients) {
                if (client.seq_id == -1 && g_seq_id < n_seq) {
                    client.seq_id = g_seq_id;
@@ -348,17 +350,17 @@ int main(int argc, char ** argv) {
                    tokens_prompt = common_tokenize(ctx, client.prompt, false);

                    for (size_t i = 0; i < tokens_prompt.size(); ++i) {
-                        batch.add(tokens_prompt[i], client.n_past++, client.id + 1, false);
+                        common_batch_add(batch, tokens_prompt[i], client.n_past++, { client.id + 1 }, false);
                    }

                    // extract the logits only for the last token
-                    if (batch.size() > 0) {
-                        batch.set_output(batch.size() - 1, true);
+                    if (batch.n_tokens > 0) {
+                        batch.logits[batch.n_tokens - 1] = true;
                    }

                    client.n_prompt  = tokens_prompt.size();
                    client.n_decoded = 0;
-                    client.i_batch   = batch.size() - 1;
+                    client.i_batch   = batch.n_tokens - 1;

                    LOG_INF("\033[31mClient %3d, seq %4d, junk = %4d, prompt = %d, started decoding ...\033[0m\n", client.id, client.seq_id, n_junk_cur, client.n_prompt);

@@ -372,7 +374,7 @@ int main(int argc, char ** argv) {
            }
        }

-        if (batch.size() == 0) {
+        if (batch.n_tokens == 0) {
            break;
        }

@@ -381,17 +383,27 @@ int main(int argc, char ** argv) {

        int32_t i_next = 0;

-        for (int32_t i = 0; i < batch.size(); i = i_next) {
+        for (int32_t i = 0; i < batch.n_tokens; i = i_next) {
            // experiment: process in powers of 2
-            //if (i + n_batch > (int32_t) batch.size() && n_batch > 32) {
+            //if (i + n_batch > (int32_t) batch.n_tokens && n_batch > 32) {
            //    n_batch /= 2;
            //    i -= n_batch;
            //    continue;
            //}

-            const int32_t n_tokens = std::min(n_batch, batch.size() - i);
+            const int32_t n_tokens = std::min(n_batch, batch.n_tokens - i);

-            const int ret = llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get_sub_batch(i, n_tokens));
+            llama_batch batch_view = {
+                n_tokens,
+                batch.token    + i,
+                nullptr,
+                batch.pos      + i,
+                batch.n_seq_id + i,
+                batch.seq_id   + i,
+                batch.logits   + i,
+            };
+
+            const int ret = llama_decode(ctx, batch_view);
            if (ret != 0) {
                if (n_batch == 1 || ret < 0) {
                    // if you get here, it means the KV cache is full - try increasing it via the context size
@@ -499,6 +511,7 @@ int main(int argc, char ** argv) {
    // TODO: print sampling/grammar timings for all clients
    llama_perf_context_print(ctx);

+    llama_batch_free(batch);

    llama_backend_free();

@@ -125,7 +125,7 @@ int main(int argc, char ** argv) {
    LOG_INF("prompt tokens: %d\n", n_tokens_all);
    //LOG_INF("prompt: %s\n", params.prompt.c_str());

-    common_batch batch(ctx);
+    llama_batch batch = llama_batch_init(params.n_batch, 0, 1);

    int n_past = 0;

@@ -144,17 +144,17 @@ int main(int argc, char ** argv) {
            n_past = llama_memory_seq_pos_max(mem, 0) + 1;
        }

-        batch.clear();
+        common_batch_clear(batch);

        for (int j = 0; j < n_batch && i + j < n_tokens_all; j++) {
-            batch.add(tokens_list[i + j], n_past++, 0, false);
+            common_batch_add(batch, tokens_list[i + j], n_past++, { 0 }, false);
        }

        if (i + n_batch >= n_tokens_all) {
-            batch.set_output(batch.size() - 1, true);
+            batch.logits[batch.n_tokens - 1] = true;
        }

-        if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()) != 0) {
+        if (llama_decode(ctx, batch) != 0) {
            LOG_INF("%s: llama_decode() failed\n", __func__);
            return 1;
        }
@@ -176,17 +176,17 @@ int main(int argc, char ** argv) {

        n_past = llama_memory_seq_pos_max(mem, 0) + 1;

-        batch.clear();
+        common_batch_clear(batch);

        for (int j = 0; j < n_batch && i + j < n_tokens_all; j++) {
-            batch.add(tokens_list[i + j], n_past++, 0, false);
+            common_batch_add(batch, tokens_list[i + j], n_past++, { 0 }, false);
        }

        if (i + n_batch >= n_tokens_all) {
-            batch.set_output(batch.size() - 1, true);
+            batch.logits[batch.n_tokens - 1] = true;
        }

-        if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()) != 0) {
+        if (llama_decode(ctx, batch) != 0) {
            LOG_ERR("%s: llama_decode() failed\n", __func__);
            return 1;
        }
@@ -223,7 +223,7 @@ int main(int argc, char ** argv) {
    while (n_cur <= n_len) {
        // sample the next token
        {
-            const llama_token new_token_id = llama_sampler_sample(smpl, ctx, batch.size() - 1);
+            const llama_token new_token_id = llama_sampler_sample(smpl, ctx, batch.n_tokens - 1);

            // is it an end of generation?
            if (llama_vocab_is_eog(vocab, new_token_id) || n_cur == n_len) {
@@ -237,16 +237,16 @@ int main(int argc, char ** argv) {
            n_decode += 1;

            // prepare the next batch
-            batch.clear();
+            common_batch_clear(batch);

            // push this new token for next evaluation
-            batch.add(new_token_id, n_past++, 0, true);
+            common_batch_add(batch, new_token_id, n_past++, { 0 }, true);
        }

        n_cur += 1;

        // evaluate the current batch with the transformer model
-        if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
+        if (llama_decode(ctx, batch)) {
            LOG_ERR("%s : failed to eval, return code %d\n", __func__, 1);
            return 1;
        }
@@ -266,6 +266,7 @@ int main(int argc, char ** argv) {

    llama_sampler_free(smpl);

+    llama_batch_free(batch);

    llama_free(ctx);
    llama_model_free(model);
@@ -75,30 +75,30 @@ static std::vector<chunk> chunk_file(const std::string & filename, int chunk_siz
    return chunks;
 }

-static void batch_add_seq(common_batch & batch, const std::vector<int32_t> & tokens, llama_seq_id seq_id) {
+static void batch_add_seq(llama_batch & batch, const std::vector<int32_t> & tokens, llama_seq_id seq_id) {
    size_t n_tokens = tokens.size();
    for (size_t i = 0; i < n_tokens; i++) {
-        batch.add(tokens[i], i, seq_id, true);
+        common_batch_add(batch, tokens[i], i, { seq_id }, true);
    }
 }

-static void batch_process(llama_context * ctx, common_batch & batch, float * output, int n_seq, int n_embd) {
+static void batch_process(llama_context * ctx, llama_batch & batch, float * output, int n_seq, int n_embd) {
    // clear previous kv_cache values (irrelevant for embeddings)
    llama_memory_clear(llama_get_memory(ctx), false);

    // run model
-    LOG_INF("%s: n_tokens = %d, n_seq = %d\n", __func__, batch.size(), n_seq);
-    if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch.get()) < 0) {
+    LOG_INF("%s: n_tokens = %d, n_seq = %d\n", __func__, batch.n_tokens, n_seq);
+    if (llama_decode(ctx, batch) < 0) {
        LOG_ERR("%s : failed to process\n", __func__);
    }

-    for (int i = 0; i < batch.size(); i++) {
-        if (!batch.tokens[i].output) {
+    for (int i = 0; i < batch.n_tokens; i++) {
+        if (!batch.logits[i]) {
            continue;
        }

        // try to get sequence embeddings - supported only when pooling_type is not NONE
-        const float * embd = llama_get_embeddings_seq(ctx, batch.tokens[i].seq_id);
+        const float * embd = llama_get_embeddings_seq(ctx, batch.seq_id[i][0]);
        if (embd == NULL) {
            embd = llama_get_embeddings_ith(ctx, i);
            if (embd == NULL) {
@@ -107,7 +107,7 @@ static void batch_process(llama_context * ctx, common_batch & batch, float * out
            }
        }

-        float * out = output + batch.tokens[i].seq_id * n_embd;
+        float * out = output + batch.seq_id[i][0] * n_embd;
        common_embd_normalize(embd, out, n_embd, 2);
    }
 }
@@ -217,7 +217,7 @@ int main(int argc, char ** argv) {

    // initialize batch
    const int n_chunks = chunks.size();
-    common_batch batch(ctx);
+    struct llama_batch batch = llama_batch_init(n_batch, 0, 1);

    // allocate output
    const int n_embd_out = llama_model_n_embd_out(model);
@@ -234,10 +234,10 @@ int main(int argc, char ** argv) {
        const uint64_t n_toks = inp.size();

        // encode if at capacity
-        if (batch.size() + n_toks > n_batch || s >= llama_n_seq_max(ctx)) {
+        if (batch.n_tokens + n_toks > n_batch || s >= llama_n_seq_max(ctx)) {
            float * out = emb + p * n_embd_out;
            batch_process(ctx, batch, out, s, n_embd_out);
-            batch.clear();
+            common_batch_clear(batch);
            p += s;
            s = 0;
        }
@@ -258,7 +258,7 @@ int main(int argc, char ** argv) {
        chunks[i].tokens.clear();
    }

-    common_batch query_batch(ctx);
+    struct llama_batch query_batch = llama_batch_init(n_batch, 0, 1);

    // start loop, receive query and return top k similar chunks based on cosine similarity
    std::string query;
@@ -272,7 +272,7 @@ int main(int argc, char ** argv) {
        std::vector<float> query_emb(n_embd_out, 0);
        batch_process(ctx, query_batch, query_emb.data(), 1, n_embd_out);

-        query_batch.clear();
+        common_batch_clear(query_batch);

        // compute cosine similarities
        {
@@ -302,5 +302,6 @@ int main(int argc, char ** argv) {
    llama_perf_context_print(ctx);

    // clean up
+    llama_batch_free(query_batch);
    llama_backend_free();
 }
@@ -6,17 +6,6 @@
 #include <string>
 #include <vector>

-// fill the batch with tokens at consecutive positions starting from pos_0, output logits only for the last one
-static void batch_set_tokens(llama_batch_ext * batch, const llama_token * tokens, int32_t n_tokens, llama_pos pos_0) {
-    llama_batch_ext_clear(batch);
-    for (int32_t i = 0; i < n_tokens; ++i) {
-        const int32_t idx = llama_batch_ext_add_token(batch, 0, tokens[i]);
-        const llama_pos pos = pos_0 + i;
-        llama_batch_ext_set_pos(batch, idx, &pos);
-    }
-    llama_batch_ext_set_output_logits(batch, n_tokens - 1, true);
-}
-
 static void print_usage(int, char ** argv) {
    printf("\nexample usage:\n");
    printf("\n    %s -m model.gguf [-c context_size] [-ngl n_gpu_layers]\n", argv[0]);
@@ -76,7 +65,8 @@ int main(int argc, char ** argv) {
        }
    }, nullptr);

-    llama_backend_init();
+    // load dynamic backends
+    ggml_backend_load_all();

    // initialize the model
    llama_model_params model_params = llama_model_default_params();
@@ -107,8 +97,6 @@ int main(int argc, char ** argv) {
    llama_sampler_chain_add(smpl, llama_sampler_init_temp(0.8f));
    llama_sampler_chain_add(smpl, llama_sampler_init_dist(LLAMA_DEFAULT_SEED));

-    llama_batch_ext * batch = llama_batch_ext_init(ctx);
-
    // helper function to evaluate a prompt and generate a response
    auto generate = [&](const std::string & prompt) {
        std::string response;
@@ -122,25 +110,20 @@ int main(int argc, char ** argv) {
            GGML_ABORT("failed to tokenize the prompt\n");
        }

-        // the tokens to evaluate next: the prompt, then the sampled token
-        const llama_token * tokens = prompt_tokens.data();
-        int n_tokens = prompt_tokens.size();
-
+        // prepare a batch for the prompt
+        llama_batch batch = llama_batch_get_one(prompt_tokens.data(), prompt_tokens.size());
        llama_token new_token_id;
        while (true) {
            // check if we have enough space in the context to evaluate this batch
            int n_ctx = llama_n_ctx(ctx);
            int n_ctx_used = llama_memory_seq_pos_max(llama_get_memory(ctx), 0) + 1;
-            if (n_ctx_used + n_tokens > n_ctx) {
+            if (n_ctx_used + batch.n_tokens > n_ctx) {
                printf("\033[0m\n");
                fprintf(stderr, "context size exceeded\n");
                exit(0);
            }

-            // positions continue from the memory
-            batch_set_tokens(batch, tokens, n_tokens, n_ctx_used);
-
-            int ret = llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch);
+            int ret = llama_decode(ctx, batch);
            if (ret != 0) {
                GGML_ABORT("failed to decode, ret = %d\n", ret);
            }
@@ -165,8 +148,7 @@ int main(int argc, char ** argv) {
            response += piece;

            // prepare the next batch with the sampled token
-            tokens   = &new_token_id;
-            n_tokens = 1;
+            batch = llama_batch_get_one(&new_token_id, 1);
        }

        return response;
@@ -220,7 +202,6 @@ int main(int argc, char ** argv) {
    for (auto & msg : messages) {
        free(const_cast<char *>(msg.content));
    }
-    llama_batch_ext_free(batch);
    llama_sampler_free(smpl);
    llama_free(ctx);
    llama_model_free(model);
@@ -5,17 +5,6 @@
 #include <string>
 #include <vector>

-// fill the batch with tokens at consecutive positions starting from pos_0, output logits only for the last one
-static void batch_set_tokens(llama_batch_ext * batch, const llama_token * tokens, int32_t n_tokens, llama_pos pos_0) {
-    llama_batch_ext_clear(batch);
-    for (int32_t i = 0; i < n_tokens; ++i) {
-        const int32_t idx = llama_batch_ext_add_token(batch, 0, tokens[i]);
-        const llama_pos pos = pos_0 + i;
-        llama_batch_ext_set_pos(batch, idx, &pos);
-    }
-    llama_batch_ext_set_output_logits(batch, n_tokens - 1, true);
-}
-
 static void print_usage(int, char ** argv) {
    printf("\nexample usage:\n");
    printf("\n    %s -m model.gguf [-n n_predict] [-ngl n_gpu_layers] [prompt]\n", argv[0]);
@@ -88,7 +77,9 @@ int main(int argc, char ** argv) {
        }
    }

-    llama_backend_init();
+    // load dynamic backends
+
+    ggml_backend_load_all();

    // initialize the model

@@ -155,13 +146,10 @@ int main(int argc, char ** argv) {

    // prepare a batch for the prompt

-    llama_batch_ext * batch = llama_batch_ext_init(ctx);
-    int n_tokens = n_prompt; // number of tokens in the current batch
-
-    batch_set_tokens(batch, prompt_tokens.data(), n_prompt, 0);
+    llama_batch batch = llama_batch_get_one(prompt_tokens.data(), prompt_tokens.size());

    if (llama_model_has_encoder(model)) {
-        if (llama_process(ctx, LLAMA_PROCESS_TYPE_ENCODE, batch)) {
+        if (llama_encode(ctx, batch)) {
            fprintf(stderr, "%s : failed to eval\n", __func__);
            return 1;
        }
@@ -171,8 +159,7 @@ int main(int argc, char ** argv) {
            decoder_start_token_id = llama_vocab_bos(vocab);
        }

-        batch_set_tokens(batch, &decoder_start_token_id, 1, 0);
-        n_tokens = 1;
+        batch = llama_batch_get_one(&decoder_start_token_id, 1);
    }

    // main loop
@@ -181,14 +168,14 @@ int main(int argc, char ** argv) {
    int n_decode = 0;
    llama_token new_token_id;

-    for (int n_pos = 0; n_pos + n_tokens < n_prompt + n_predict; ) {
+    for (int n_pos = 0; n_pos + batch.n_tokens < n_prompt + n_predict; ) {
        // evaluate the current batch with the transformer model
-        if (llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, batch)) {
+        if (llama_decode(ctx, batch)) {
            fprintf(stderr, "%s : failed to eval, return code %d\n", __func__, 1);
            return 1;
        }

-        n_pos += n_tokens;
+        n_pos += batch.n_tokens;

        // sample the next token
        {
@@ -210,8 +197,7 @@ int main(int argc, char ** argv) {
            fflush(stdout);

            // prepare the next batch with the sampled token
-            batch_set_tokens(batch, &new_token_id, 1, n_pos);
-            n_tokens = 1;
+            batch = llama_batch_get_one(&new_token_id, 1);

            n_decode += 1;
        }
@@ -229,7 +215,6 @@ int main(int argc, char ** argv) {
    llama_perf_context_print(ctx);
    fprintf(stderr, "\n");

-    llama_batch_ext_free(batch);
    llama_sampler_free(smpl);
    llama_free(ctx);
    llama_model_free(model);
@@ -125,12 +125,12 @@ int main(int argc, char ** argv) {

    // eval the prompt on the target and feed it to the speculative implementation(s)
    {
-        common_batch batch_prompt(ctx_tgt);
+        llama_batch batch_prompt = llama_batch_init(inp.size(), 0, 1);
        for (size_t i = 0; i < inp.size() - 1; ++i) {
-            batch_prompt.add(inp[i], i, seq_id, false);
+            common_batch_add(batch_prompt, inp[i], i, { seq_id }, false);
        }

-        llama_process(ctx_tgt, LLAMA_PROCESS_TYPE_DECODE, batch_prompt.get());
+        llama_decode(ctx_tgt, batch_prompt);

        if (!common_speculative_process(spec, batch_prompt)) {
            LOG_ERR("%s", "failed to process speculative prompt\n");
@@ -149,7 +149,7 @@ int main(int argc, char ** argv) {

    common_speculative_begin(spec, seq_id, prompt_tgt);

-    common_batch batch_tgt(ctx_tgt);
+    llama_batch batch_tgt = llama_batch_init(llama_n_batch(ctx_tgt), 0, 1);

    llama_tokens draft;

@@ -219,17 +219,18 @@ int main(int argc, char ** argv) {
        }

        // always have a token to evaluate from before - id_last
-        batch_tgt.clear();
-        batch_tgt.add(id_last, n_past++, seq_id, true);
+        common_batch_clear(batch_tgt);
+        common_batch_add  (batch_tgt, id_last, n_past++, { seq_id }, true);

        // evaluate the target model on [id_last, draft0, draft1, ..., draftN-1]
        {
            for (size_t i = 0; i < draft.size(); ++i) {
-                batch_tgt.add(draft[i], n_past + i, seq_id, true);
+                common_batch_add(batch_tgt, draft[i], n_past + i, { seq_id }, true);
            }

+            //LOG_DBG("target batch: %s\n", string_from(ctx_tgt, batch_tgt).c_str());

-            llama_process(ctx_tgt, LLAMA_PROCESS_TYPE_DECODE, batch_tgt.get());
+            llama_decode(ctx_tgt, batch_tgt);
        }

        // feed the batch to the speculative implementation(s) - this drives the draft model, MTP, Eagle3, etc.
@@ -364,6 +365,7 @@ int main(int argc, char ** argv) {
    LOG_INF("target:\n\n");
    common_perf_print(ctx_tgt, smpl.get());

+    llama_batch_free(batch_tgt);

    common_speculative_free(spec);

@@ -190,16 +190,9 @@ int main(int argc, char ** argv) {
    const auto t_enc_start = ggml_time_us();

    // eval the prompt with both models
-    {
-        common_batch batch = common_batch_get_one(ctx_tgt, inp.data(), n_input - 1);
-        llama_process(ctx_tgt, LLAMA_PROCESS_TYPE_DECODE, batch.get());
-
-        batch = common_batch_get_one(ctx_tgt, &inp.back(), 1);
-        llama_process(ctx_tgt, LLAMA_PROCESS_TYPE_DECODE, batch.get());
-
-        batch = common_batch_get_one(ctx_dft, inp.data(), n_input);
-        llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch.get());
-    }
+    llama_decode(ctx_tgt, llama_batch_get_one( inp.data(), n_input - 1));
+    llama_decode(ctx_tgt, llama_batch_get_one(&inp.back(),           1));
+    llama_decode(ctx_dft, llama_batch_get_one( inp.data(), n_input));

    const auto t_enc_end = ggml_time_us();

@@ -230,8 +223,8 @@ int main(int argc, char ** argv) {
        drafts[s].smpl = common_sampler_init(model_dft, params.sampling);
    }

-    common_batch batch_dft(ctx_dft);
-    common_batch batch_tgt(ctx_tgt);
+    llama_batch batch_dft = llama_batch_init(llama_n_batch(ctx_dft), 0, 1);
+    llama_batch batch_tgt = llama_batch_init(llama_n_batch(ctx_tgt), 0, n_seq_dft);

    const auto t_dec_start = ggml_time_us();

@@ -472,12 +465,12 @@ int main(int argc, char ** argv) {
            drafts[0].dists.push_back(std::vector<llama_token_data>());
            drafts[0].i_batch_tgt.push_back(0);

-            batch_dft.clear();
-            batch_dft.add(token_id, n_past_dft, 0, true);
+            common_batch_clear(batch_dft);
+            common_batch_add  (batch_dft, token_id, n_past_dft, { 0 }, true);

            llama_memory_seq_rm(mem_dft, 0, n_past_dft, -1);
            // LOG_DBG("dft batch: %s\n", LOG_BATCH_TOSTR_PRETTY(ctx_dft, batch_dft).c_str());
-            llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch_dft.get());
+            llama_decode(ctx_dft, batch_dft);

            ++n_past_dft;
        }
@@ -502,12 +495,12 @@ int main(int argc, char ** argv) {
        drafts[0].drafting    = true;
        drafts[0].i_batch_dft = 0;

-        batch_tgt.clear();
-        batch_tgt.add(drafts[0].tokens[0], n_past_tgt, 0, true);
+        common_batch_clear(batch_tgt);
+        common_batch_add  (batch_tgt, drafts[0].tokens[0], n_past_tgt, { 0 }, true);

        // sample n_draft tokens from the draft model using tree-based sampling
        for (int i = 0; i < n_draft; ++i) {
-            batch_dft.clear();
+            batch_dft.n_tokens = 0;

            for (int s = 0; s < n_seq_dft; ++s) {
                drafts[s].skip = false;
@@ -538,8 +531,14 @@ int main(int argc, char ** argv) {
                        llama_memory_seq_cp(mem_dft, s, n_seq_cur, -1, -1);

                        // all previous tokens from this branch are now also part of the new branch
-                        for (int t : drafts[s].i_batch_tgt) {
-                            batch_tgt.add_seq(t, n_seq_cur);
+                        for (int t = 0; t < batch_tgt.n_tokens; ++t) {
+                            for (int p = 0; p < batch_tgt.n_seq_id[t]; ++p) {
+                                if (batch_tgt.seq_id[t][p] == s) {
+                                    batch_tgt.seq_id[t][batch_tgt.n_seq_id[t]] = n_seq_cur;
+                                    batch_tgt.n_seq_id[t]++;
+                                    break;
+                                }
+                            }
                        }

                        // copy the draft state
@@ -578,32 +577,32 @@ int main(int argc, char ** argv) {
                    drafts[s].dists.push_back({cur_p->data, cur_p->data + cur_p->size});

                    // add unique drafted tokens to the target batch
-                    drafts[s].i_batch_tgt.push_back(batch_tgt.size());
+                    drafts[s].i_batch_tgt.push_back(batch_tgt.n_tokens);

-                    batch_tgt.add(id, n_past_tgt + i + 1, s, true);
+                    common_batch_add(batch_tgt, id, n_past_tgt + i + 1, { s }, true);

                    // add the token to the batch for batched decoding with the draft model
-                    drafts[s].i_batch_dft = batch_dft.size();
+                    drafts[s].i_batch_dft = batch_dft.n_tokens;

-                    batch_dft.add(id, n_past_cur, s, true);
+                    common_batch_add(batch_dft, id, n_past_cur, { s }, true);

-                    if (batch_tgt.size() > n_draft) {
+                    if (batch_tgt.n_tokens > n_draft) {
                        drafts[s].drafting = false;
                    }
                }
            }

            // no sequence is drafting anymore
-            if (batch_dft.size() == 0) {
+            if (batch_dft.n_tokens == 0) {
                break;
            }

            // evaluate the drafted tokens on the draft model
-            llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch_dft.get());
+            llama_decode(ctx_dft, batch_dft);
            ++n_past_cur;
            ++n_drafted;

-            if (batch_tgt.size() > n_draft) {
+            if (batch_tgt.n_tokens > n_draft) {
                break;
            }
        }
@@ -616,7 +615,7 @@ int main(int argc, char ** argv) {
            }

            // LOG_DBG("target batch: %s\n", LOG_BATCH_TOSTR_PRETTY(ctx_tgt, batch_tgt).c_str());
-            llama_process(ctx_tgt, LLAMA_PROCESS_TYPE_DECODE, batch_tgt.get());
+            llama_decode(ctx_tgt, batch_tgt);
            ++n_past_tgt;
        }

@@ -659,6 +658,7 @@ int main(int argc, char ** argv) {
        common_sampler_free(drafts[s].smpl);
    }

+    llama_batch_free(batch_dft);

    llama_backend_free();

@@ -197,7 +197,6 @@ set(GGML_BLAS_VENDOR ${GGML_BLAS_VENDOR_DEFAULT} CACHE STRING
 option(GGML_LLAMAFILE                       "ggml: use LLAMAFILE"                             ${GGML_LLAMAFILE_DEFAULT})

 option(GGML_CUDA                            "ggml: use CUDA"                                  OFF)
-set   (GGML_CUDA_CCCL_VERSION "" CACHE STRING "ggml: CCCL git tag to fetch, empty to use the version bundled with the installed CUDA Toolkit")
 option(GGML_MUSA                            "ggml: use MUSA"                                  OFF)
 option(GGML_CUDA_FORCE_MMQ                  "ggml: use mmq kernels instead of cuBLAS"         OFF)
 option(GGML_CUDA_FORCE_CUBLAS               "ggml: always use cuBLAS instead of mmq kernels"  OFF)
@@ -10,8 +10,6 @@ extern "C" {
 // device buffer
 GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_zdnn_buffer_type(void);

-GGML_BACKEND_API bool ggml_backend_is_zdnn(ggml_backend_t backend);
-
 GGML_BACKEND_API ggml_backend_reg_t ggml_backend_zdnn_reg(void);

 #ifdef __cplusplus
@@ -2331,22 +2331,20 @@ static enum ggml_status ggml_backend_meta_graph_compute(ggml_backend_t backend,
            if (node->flags & GGML_TENSOR_FLAG_COMPUTE) {
                continue;
            }
-            if (ggml_nelements(node) > 0) {
-                ggml_tensor * node_zero = get_node_aux(node);
-                node_zero->op = GGML_OP_FILL;
-                node_zero->src[0] = node; // only used for the shape, the data is not read
-                ggml_set_op_params_f32(node_zero, 0, 0.0f);
-                node_zero->data = node->data;
-                node_zero->buffer = node->buffer;
-                node_zero->flags |= GGML_TENSOR_FLAG_COMPUTE;
+            ggml_tensor * node_zero = get_node_aux(node);
+            node_zero->op = GGML_OP_SCALE; // FIXME 0.0f * NaN == NaN
+            node_zero->src[0] = node;
+            ggml_set_op_params_f32(node_zero, 0, 0.0f);
+            node_zero->data = node->data;
+            node_zero->buffer = node->buffer;
+            node_zero->flags |= GGML_TENSOR_FLAG_COMPUTE;

-                step_cgraphs[j] = get_cgraph_aux();
-                step_cgraphs[j]->nodes[0] = node_zero;
-                step_cgraphs[j]->n_nodes = 1;
-                const ggml_status status = ggml_backend_graph_compute_async(bcj.backend, step_cgraphs[j]);
-                if (status != GGML_STATUS_SUCCESS) {
-                    return status;
-                }
+            step_cgraphs[j] = get_cgraph_aux();
+            step_cgraphs[j]->nodes[0] = node_zero;
+            step_cgraphs[j]->n_nodes = 1;
+            const ggml_status status = ggml_backend_graph_compute_async(bcj.backend, step_cgraphs[j]);
+            if (status != GGML_STATUS_SUCCESS) {
+                return status;
            }
        }
        std::fill(step_cgraphs.begin(), step_cgraphs.end(), nullptr);
@@ -1372,6 +1372,30 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra
                const int src_backend_id = sched->hv_tensor_backend_ids[src_id];
                GGML_ASSERT(src_backend_id != -1); // all inputs should be assigned by now

+                if (src->flags & GGML_TENSOR_FLAG_INPUT && sched->n_copies > 1) {
+                    if (tensor_id_copy(src_id, src_backend_id, 0) == NULL) {
+                        ggml_backend_t backend = sched->backends[src_backend_id];
+                        for (int c = 0; c < sched->n_copies; c++) {
+                            struct ggml_tensor * tensor_copy;
+                            if (c == sched->cur_copy) {
+                                tensor_copy = src; // use the original tensor as the current copy
+                            } else {
+                                tensor_copy = ggml_dup_tensor_layout(sched->ctx, src);
+                                ggml_format_name(tensor_copy, "%s#%s#%d", ggml_backend_name(backend), src->name, c);
+                            }
+                            ggml_set_input(tensor_copy);
+                            ggml_set_output(tensor_copy); // prevent ggml-alloc from overwriting the tensor
+                            tensor_id_copy(src_id, src_backend_id, c) = tensor_copy;
+                            SET_CAUSE(tensor_copy, "4.cpy");
+                        }
+                        int n_graph_inputs = sched->n_graph_inputs++;
+                        if (n_graph_inputs >= sched->graph_inputs_capacity) {
+                            ggml_backend_sched_graph_inputs_grow(sched);
+                        }
+                        sched->graph_inputs[n_graph_inputs] = src;
+                    }
+                }
+
                if (src_backend_id != cur_backend_id && !ggml_backend_sched_buffer_supported(sched, src, cur_backend_id)) {
                    // create a copy of the input in the split's backend
                    if (tensor_id_copy(src_id, cur_backend_id, 0) == NULL) {
@@ -1404,46 +1428,6 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra
        ggml_backend_sched_print_assignments(sched, graph);
    }

-    // pass 6: collect all input tensors into graph_inputs
-    //         this includes inputs not consumed by any node (e.g. the embeddings input of a text-only batch) so that
-    //         the graph composition does not depend on which inputs are used, which would otherwise cause graph
-    //         reallocations when switching between different types of batches [GGML_SCHED_DEBUG_REALLOC]
-    if (sched->n_copies > 1) {
-        for (int i = 0; i < graph->n_leafs; i++) {
-            struct ggml_tensor * leaf = graph->leafs[i];
-            if ((leaf->flags & GGML_TENSOR_FLAG_INPUT) == 0) {
-                continue;
-            }
-
-            const size_t leaf_id = hash_id(leaf);
-            const int leaf_backend_id = tensor_backend_id(leaf);
-            GGML_ASSERT(leaf_backend_id != -1); // all leafs should be assigned by now
-
-            if (tensor_id_copy(leaf_id, leaf_backend_id, 0) == NULL) {
-                ggml_backend_t backend = sched->backends[leaf_backend_id];
-                for (int c = 0; c < sched->n_copies; c++) {
-                    struct ggml_tensor * tensor_copy;
-                    if (c == sched->cur_copy) {
-                        tensor_copy = leaf; // use the original tensor as the current copy
-                    } else {
-                        tensor_copy = ggml_dup_tensor_layout(sched->ctx, leaf);
-                        ggml_format_name(tensor_copy, "%s#%s#%d", ggml_backend_name(backend), leaf->name, c);
-                    }
-                    ggml_set_input(tensor_copy);
-                    ggml_set_output(tensor_copy); // prevent ggml-alloc from overwriting the tensor
-                    tensor_id_copy(leaf_id, leaf_backend_id, c) = tensor_copy;
-                    SET_CAUSE(tensor_copy, "6.cpy");
-                }
-            }
-
-            int n_graph_inputs = sched->n_graph_inputs++;
-            if (n_graph_inputs >= sched->graph_inputs_capacity) {
-                ggml_backend_sched_graph_inputs_grow(sched);
-            }
-            sched->graph_inputs[n_graph_inputs] = leaf;
-        }
-    }
-
    // swap node_backend_ids and leaf _backend_ids with prevs
    {
        int * tmp = sched->node_backend_ids;
@@ -84,17 +84,10 @@ if (BLAS_FOUND)
        add_compile_definitions(GGML_BLAS_USE_OPENBLAS)
    endif()

-    # Vendors with the BLIS API take GGML_BLAS_USE_BLIS (blis.h, bli_thread_set_num_threads).
-    # AOCL also takes GGML_BLAS_USE_AOCL so its device label is AOCL-BLAS, not BLIS.
-    # Neither flag selects the linked library; FindBLAS does that from GGML_BLAS_VENDOR.
    if ("${GGML_BLAS_VENDOR}" MATCHES "FLAME" OR "${GGML_BLAS_VENDOR}" MATCHES "AOCL" OR "${GGML_BLAS_VENDOR}" MATCHES "AOCL_mt")
        add_compile_definitions(GGML_BLAS_USE_BLIS)
    endif()

-    if ("${GGML_BLAS_VENDOR}" MATCHES "AOCL" OR "${GGML_BLAS_VENDOR}" MATCHES "AOCL_mt")
-        add_compile_definitions(GGML_BLAS_USE_AOCL)
-    endif()
-
    if ("${GGML_BLAS_VENDOR}" MATCHES "NVPL")
        add_compile_definitions(GGML_BLAS_USE_NVPL)
    endif()
@@ -330,8 +330,6 @@ static const char * ggml_backend_blas_device_get_description(ggml_backend_dev_t
        return "Accelerate";
    #elif defined(GGML_BLAS_USE_MKL)
        return "MKL";
-    #elif defined(GGML_BLAS_USE_AOCL)
-        return "AOCL-BLAS";
    #elif defined(GGML_BLAS_USE_BLIS)
        return "BLIS";
    #elif defined(GGML_BLAS_USE_NVPL)
@@ -1334,11 +1334,9 @@ UseGgmlGemm1:;
        const size_t nbw3 = nbw2*ne12;

        assert(params->wsize >= ne13*nbw3);
-        // src1 is either packed from F32 into vec_dot_type, or widened from F16 or BF16 into the F32 work buffer
-        const bool widen = src1->type != GGML_TYPE_F32;
-
-        GGML_ASSERT(!widen || vec_dot_type == GGML_TYPE_F32);
-        GGML_ASSERT(!widen || src1->type == GGML_TYPE_F16 || src1->type == GGML_TYPE_BF16);
+        GGML_ASSERT(src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16);
+        // the F16 path below writes plain floats into wdata, so it needs an F32 vec_dot_type
+        GGML_ASSERT(src1->type == GGML_TYPE_F32 || vec_dot_type == GGML_TYPE_F32);

    #if 0
        for (int64_t i13 = 0; i13 < ne13; ++i13) {
@@ -1351,24 +1349,20 @@ UseGgmlGemm1:;
            }
        }
    #else
-        const int64_t bs = ggml_blck_size(vec_dot_type);
-
        for (int64_t i13 = 0; i13 < ne13; ++i13) {
            for (int64_t i12 = 0; i12 < ne12; ++i12) {
                for (int64_t i11 = 0; i11 < ne11; ++i11) {
+                    size_t bs = ggml_blck_size(vec_dot_type);
                    int64_t ne10_block_start = (ith * ne10/bs) / nth;
                    int64_t ne10_block_end   = ((ith + 1) * ne10/bs) / nth;
-                    const void * src1_row  = (const char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11 + ne10_block_start*bs*nb10;
-                    void       * wdata_row =                    wdata  + i13*nbw3 + i12*nbw2 + i11*nbw1 + ne10_block_start*nbw0;
+                    const char * src1_block = (const char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11 + ne10_block_start*bs*nb10;
+                    char * dst_block = wdata + i13*nbw3 + i12*nbw2 + i11*nbw1 + ne10_block_start*nbw0;
+                    const int64_t n_block = (ne10_block_end - ne10_block_start) * bs;

-                    const int64_t ne10_block_size = (ne10_block_end - ne10_block_start) * bs;
-
-                    if (src1->type == GGML_TYPE_F16) {
-                        ggml_cpu_fp16_to_fp32((const ggml_fp16_t *) src1_row, (float *) wdata_row, ne10_block_size);
-                    } else if (src1->type == GGML_TYPE_BF16) {
-                        ggml_cpu_bf16_to_fp32((const ggml_bf16_t *) src1_row, (float *) wdata_row, ne10_block_size);
+                    if (src1->type == GGML_TYPE_F32) {
+                        from_float((const float *) src1_block, dst_block, n_block);
                    } else {
-                        from_float((const float *) src1_row, wdata_row, ne10_block_size);
+                        ggml_cpu_fp16_to_fp32((const ggml_fp16_t *) src1_block, (float *) dst_block, n_block);
                    }
                }
            }
@@ -455,10 +455,7 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st
                src0->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32) {
                return src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16;
            }
-            // BF16 in src1 is widened into the F32 work buffer
-            return src1->type == GGML_TYPE_F32 ||
-                   src1->type == ggml_get_type_traits_cpu(src0->type)->vec_dot_type ||
-                   (src1->type == GGML_TYPE_BF16 && ggml_get_type_traits_cpu(src0->type)->vec_dot_type == GGML_TYPE_F32);
+            return src1->type == GGML_TYPE_F32 || src1->type == ggml_get_type_traits_cpu(src0->type)->vec_dot_type;
        case GGML_OP_SOFT_MAX_BACK: {
            if (op->src[0]->type != GGML_TYPE_F32 || op->src[1]->type != GGML_TYPE_F32) {
                return false;
@@ -2106,46 +2106,6 @@ void ggml_compute_forward_concat(
    }
 }

-// row loop shared by the BF16 unary ops
-template <void (*vec_op)(const int, ggml_bf16_t *, const ggml_bf16_t *)>
-static void ggml_compute_forward_unary_bf16_rows(
-        const ggml_compute_params * params,
-        ggml_tensor * dst) {
-
-    const ggml_tensor * src0 = dst->src[0];
-
-    assert(ggml_is_contiguous_rows(src0));
-    assert(ggml_are_same_shape(src0, dst));
-
-    GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne)
-    GGML_TENSOR_LOCALS(size_t,  nb0, src0, nb)
-    GGML_TENSOR_LOCALS(int64_t, ne,  dst,  ne)
-    GGML_TENSOR_LOCALS(size_t,  nb,  dst,  nb)
-
-    const int ith = params->ith;
-    const int nth = params->nth;
-
-    const int nc = src0->ne[0];
-    const int nr = ggml_nrows(src0);
-
-    // rows per thread
-    const int dr = (nr + nth - 1)/nth;
-
-    // row range for this thread
-    const int ir0 = dr*ith;
-    const int ir1 = MIN(ir0 + dr, nr);
-
-    for (int ir = ir0; ir < ir1; ++ir) {
-        const int i3 = ir/(ne02*ne01);
-        const int i2 = (ir - i3*ne02*ne01)/ne01;
-        const int i1 = (ir - i3*ne02*ne01 - i2*ne01);
-
-        vec_op(nc,
-                (ggml_bf16_t *) ((char *) dst->data  + i3*nb3  + i2*nb2  + i1*nb1),
-                (ggml_bf16_t *) ((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01));
-    }
-}
-
 // ggml_compute_forward_gelu

 static void ggml_compute_forward_gelu_f32(
@@ -2243,12 +2203,6 @@ static void ggml_compute_forward_gelu_f16(
    }
 }

-static void ggml_compute_forward_gelu_bf16(
-        const ggml_compute_params * params,
-        ggml_tensor * dst) {
-    ggml_compute_forward_unary_bf16_rows<ggml_vec_gelu_bf16>(params, dst);
-}
-
 static void ggml_compute_forward_gelu(
        const ggml_compute_params * params,
        ggml_tensor * dst) {
@@ -2264,10 +2218,6 @@ static void ggml_compute_forward_gelu(
            {
                ggml_compute_forward_gelu_f16(params, dst);
            } break;
-        case GGML_TYPE_BF16:
-            {
-                ggml_compute_forward_gelu_bf16(params, dst);
-            } break;
        default:
            {
                GGML_ABORT("fatal error");
@@ -2483,12 +2433,6 @@ static void ggml_compute_forward_gelu_erf_f16(
    }
 }

-static void ggml_compute_forward_gelu_erf_bf16(
-        const ggml_compute_params * params,
-        ggml_tensor * dst) {
-    ggml_compute_forward_unary_bf16_rows<ggml_vec_gelu_erf_bf16>(params, dst);
-}
-
 static void ggml_compute_forward_gelu_erf(
        const ggml_compute_params * params,
        ggml_tensor * dst) {
@@ -2504,10 +2448,6 @@ static void ggml_compute_forward_gelu_erf(
            {
                ggml_compute_forward_gelu_erf_f16(params, dst);
            } break;
-        case GGML_TYPE_BF16:
-            {
-                ggml_compute_forward_gelu_erf_bf16(params, dst);
-            } break;
        default:
            {
                GGML_ABORT("fatal error");
@@ -2612,12 +2552,6 @@ static void ggml_compute_forward_gelu_quick_f16(
    }
 }

-static void ggml_compute_forward_gelu_quick_bf16(
-        const ggml_compute_params * params,
-        ggml_tensor * dst) {
-    ggml_compute_forward_unary_bf16_rows<ggml_vec_gelu_quick_bf16>(params, dst);
-}
-
 static void ggml_compute_forward_gelu_quick(
        const ggml_compute_params * params,
        ggml_tensor * dst) {
@@ -2633,10 +2567,6 @@ static void ggml_compute_forward_gelu_quick(
            {
                ggml_compute_forward_gelu_quick_f16(params, dst);
            } break;
-        case GGML_TYPE_BF16:
-            {
-                ggml_compute_forward_gelu_quick_bf16(params, dst);
-            } break;
        default:
            {
                GGML_ABORT("fatal error");
@@ -2741,12 +2671,6 @@ static void ggml_compute_forward_silu_f16(
    }
 }

-static void ggml_compute_forward_silu_bf16(
-        const ggml_compute_params * params,
-        ggml_tensor * dst) {
-    ggml_compute_forward_unary_bf16_rows<ggml_vec_silu_bf16>(params, dst);
-}
-
 static void ggml_compute_forward_silu(
        const ggml_compute_params * params,
        ggml_tensor * dst) {
@@ -2762,10 +2686,6 @@ static void ggml_compute_forward_silu(
            {
                ggml_compute_forward_silu_f16(params, dst);
            } break;
-        case GGML_TYPE_BF16:
-            {
-                ggml_compute_forward_silu_bf16(params, dst);
-            } break;
        default:
            {
                GGML_ABORT("fatal error");
@@ -2969,59 +2889,6 @@ void ggml_compute_forward_silu_back(

 // ggml_compute_forward_reglu

-// row loop shared by the BF16 GLU ops
-template <void (*vec_op)(const int, ggml_bf16_t *, const ggml_bf16_t *, const ggml_bf16_t *)>
-static void ggml_compute_forward_glu_bf16_rows(
-        const ggml_compute_params * params,
-        ggml_tensor * dst) {
-
-    const ggml_tensor * src0 = dst->src[0];
-    const ggml_tensor * src1 = dst->src[1];
-    char * src0_d = (char *) src0->data;
-    char * src1_d = (char *) (src1 ? src1->data : src0->data);
-    const size_t src0_o = src0->nb[1];
-    const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1];
-
-    GGML_ASSERT(ggml_is_contiguous_1(src0));
-    GGML_ASSERT(ggml_is_contiguous_1(dst));
-
-    if (src1) {
-        GGML_ASSERT(ggml_is_contiguous_1(src1));
-        GGML_ASSERT(src0->type == src1->type);
-    }
-
-    const int ith = params->ith;
-    const int nth = params->nth;
-
-    const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2;
-    const int nr = ggml_nrows(src0);
-
-    GGML_ASSERT(dst->ne[0] == nc);
-    GGML_ASSERT(ggml_nrows(dst) == nr);
-
-    const int32_t swapped = ggml_get_op_params_i32(dst, 1);
-
-    // rows per thread
-    const int dr = (nr + nth - 1)/nth;
-
-    // row range for this thread
-    const int ir0 = dr*ith;
-    const int ir1 = MIN(ir0 + dr, nr);
-
-    for (int i1 = ir0; i1 < ir1; i1++) {
-        const ggml_bf16_t * src0_p = (const ggml_bf16_t *) (src0_d + i1*src0_o);
-        const ggml_bf16_t * src1_p = (const ggml_bf16_t *) (src1_d + i1*src1_o);
-        ggml_bf16_t       * dst_p  = (ggml_bf16_t *) ((char *) dst->data + i1*(dst->nb[1]));
-
-        if (!src1) {
-            src0_p += swapped ? nc : 0;
-            src1_p += swapped ? 0 : nc;
-        }
-
-        vec_op(nc, dst_p, src0_p, src1_p);
-    }
-}
-
 static void ggml_compute_forward_reglu_f32(
        const ggml_compute_params * params,
        ggml_tensor * dst) {
@@ -3141,12 +3008,6 @@ static void ggml_compute_forward_reglu_f16(
    }
 }

-static void ggml_compute_forward_reglu_bf16(
-        const ggml_compute_params * params,
-        ggml_tensor * dst) {
-    ggml_compute_forward_glu_bf16_rows<ggml_vec_reglu_bf16>(params, dst);
-}
-
 static void ggml_compute_forward_reglu(
        const ggml_compute_params * params,
        ggml_tensor * dst) {
@@ -3162,10 +3023,6 @@ static void ggml_compute_forward_reglu(
            {
                ggml_compute_forward_reglu_f16(params, dst);
            } break;
-        case GGML_TYPE_BF16:
-            {
-                ggml_compute_forward_reglu_bf16(params, dst);
-            } break;
        default:
            {
                GGML_ABORT("fatal error");
@@ -3294,12 +3151,6 @@ static void ggml_compute_forward_geglu_f16(
    }
 }

-static void ggml_compute_forward_geglu_bf16(
-        const ggml_compute_params * params,
-        ggml_tensor * dst) {
-    ggml_compute_forward_glu_bf16_rows<ggml_vec_geglu_bf16>(params, dst);
-}
-
 static void ggml_compute_forward_geglu(
        const ggml_compute_params * params,
        ggml_tensor * dst) {
@@ -3315,10 +3166,6 @@ static void ggml_compute_forward_geglu(
            {
                ggml_compute_forward_geglu_f16(params, dst);
            } break;
-        case GGML_TYPE_BF16:
-            {
-                ggml_compute_forward_geglu_bf16(params, dst);
-            } break;
        default:
            {
                GGML_ABORT("fatal error");
@@ -3447,12 +3294,6 @@ static void ggml_compute_forward_swiglu_f16(
    }
 }

-static void ggml_compute_forward_swiglu_bf16(
-        const ggml_compute_params * params,
-        ggml_tensor * dst) {
-    ggml_compute_forward_glu_bf16_rows<ggml_vec_swiglu_bf16>(params, dst);
-}
-
 static void ggml_compute_forward_swiglu(
        const ggml_compute_params * params,
        ggml_tensor * dst) {
@@ -3468,10 +3309,6 @@ static void ggml_compute_forward_swiglu(
            {
                ggml_compute_forward_swiglu_f16(params, dst);
            } break;
-        case GGML_TYPE_BF16:
-            {
-                ggml_compute_forward_swiglu_bf16(params, dst);
-            } break;
        default:
            {
                GGML_ABORT("fatal error");
@@ -3820,12 +3657,6 @@ static void ggml_compute_forward_geglu_erf_f16(
    }
 }

-static void ggml_compute_forward_geglu_erf_bf16(
-        const ggml_compute_params * params,
-        ggml_tensor * dst) {
-    ggml_compute_forward_glu_bf16_rows<ggml_vec_geglu_erf_bf16>(params, dst);
-}
-
 static void ggml_compute_forward_geglu_erf(
        const ggml_compute_params * params,
        ggml_tensor * dst) {
@@ -3841,10 +3672,6 @@ static void ggml_compute_forward_geglu_erf(
            {
                ggml_compute_forward_geglu_erf_f16(params, dst);
            } break;
-        case GGML_TYPE_BF16:
-            {
-                ggml_compute_forward_geglu_erf_bf16(params, dst);
-            } break;
        default:
            {
                GGML_ABORT("fatal error");
@@ -3973,12 +3800,6 @@ static void ggml_compute_forward_geglu_quick_f16(
    }
 }

-static void ggml_compute_forward_geglu_quick_bf16(
-        const ggml_compute_params * params,
-        ggml_tensor * dst) {
-    ggml_compute_forward_glu_bf16_rows<ggml_vec_geglu_quick_bf16>(params, dst);
-}
-
 static void ggml_compute_forward_geglu_quick(
        const ggml_compute_params * params,
        ggml_tensor * dst) {
@@ -3994,10 +3815,6 @@ static void ggml_compute_forward_geglu_quick(
            {
                ggml_compute_forward_geglu_quick_f16(params, dst);
            } break;
-        case GGML_TYPE_BF16:
-            {
-                ggml_compute_forward_geglu_quick_bf16(params, dst);
-            } break;
        default:
            {
                GGML_ABORT("fatal error");
@@ -4929,44 +4746,6 @@ static void ggml_compute_forward_scale_f32(
    }
 }

-static void ggml_compute_forward_scale_bf16(
-        const ggml_compute_params * params,
-        ggml_tensor * dst) {
-
-    const ggml_tensor * src0 = dst->src[0];
-
-    GGML_ASSERT(ggml_is_contiguous(src0));
-    GGML_ASSERT(ggml_is_contiguous(dst));
-    GGML_ASSERT(ggml_are_same_shape(src0, dst));
-
-    float s; // scale factor
-    float b; // bias
-
-    memcpy(&s, (float *) dst->op_params + 0, sizeof(float));
-    memcpy(&b, (float *) dst->op_params + 1, sizeof(float));
-
-    const int ith = params->ith;
-    const int nth = params->nth;
-
-    const int nc = src0->ne[0];
-    const int nr = ggml_nrows(src0);
-
-    // rows per thread
-    const int dr = (nr + nth - 1)/nth;
-
-    // row range for this thread
-    const int ir0 = dr*ith;
-    const int ir1 = MIN(ir0 + dr, nr);
-
-    for (int i1 = ir0; i1 < ir1; i1++) {
-        const ggml_bf16_t * x = (const ggml_bf16_t *) ((const char *) src0->data + i1*src0->nb[1]);
-        ggml_bf16_t       * y = (ggml_bf16_t *)       ((char *)       dst->data  + i1*dst->nb[1]);
-        for (int k = 0; k < nc; k++) {
-            y[k] = GGML_FP32_TO_BF16(s*GGML_BF16_TO_FP32(x[k]) + b);
-        }
-    }
-}
-
 void ggml_compute_forward_scale(
        const ggml_compute_params * params,
        ggml_tensor * dst) {
@@ -4978,10 +4757,6 @@ void ggml_compute_forward_scale(
            {
                ggml_compute_forward_scale_f32(params, dst);
            } break;
-        case GGML_TYPE_BF16:
-            {
-                ggml_compute_forward_scale_bf16(params, dst);
-            } break;
        default:
            {
                GGML_ABORT("fatal error");
@@ -5576,8 +5351,6 @@ static void ggml_compute_forward_get_rows_back_f32_f16(
    for (int i = 0; i < nr; ++i) {
        const int r = ((int32_t *) src1->data)[i];

-        GGML_ASSERT(r >= 0 && r < dst->ne[1]);
-
        for (int j = 0; j < nc; ++j) {
            ggml_fp16_t v = ((ggml_fp16_t *) ((char *) src0->data + i*src0->nb[1]))[j];
            ((float *) ((char *) dst->data + r*dst->nb[1]))[j] += GGML_CPU_FP16_TO_FP32(v);
@@ -5611,8 +5384,6 @@ static void ggml_compute_forward_get_rows_back_f32(
    for (int i = 0; i < nr; ++i) {
        const int r = ((int32_t *) src1->data)[i];

-        GGML_ASSERT(r >= 0 && r < dst->ne[1]);
-
        ggml_vec_add_f32(nc,
                (float *) ((char *)  dst->data + r*dst->nb[1]),
                (float *) ((char *)  dst->data + r*dst->nb[1]),
@@ -9266,11 +9037,6 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
            simd_gemm(KQ, (const float *)Q_q, K_f32, Q_TILE_SZ, DK, KV_TILE_SZ);
            ggml_vec_scale_f32(Q_TILE_SZ * KV_TILE_SZ, KQ, scale);

-            if (logit_softcap != 0.0f) {
-                ggml_vec_tanh_f32(Q_TILE_SZ * KV_TILE_SZ, KQ, KQ);
-                ggml_vec_scale_f32(Q_TILE_SZ * KV_TILE_SZ, KQ, logit_softcap);
-            }
-
            // Set padded KQ entries to -inf so softmax gives them zero weight
            if (kv_tile < KV_TILE_SZ) {
                for (int tq = 0; tq < Q_TILE_SZ; tq++) {
@@ -9280,6 +9046,11 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
                }
            }

+            if (logit_softcap != 0.0f) {
+                ggml_vec_tanh_f32(Q_TILE_SZ * KV_TILE_SZ, KQ, KQ);
+                ggml_vec_scale_f32(Q_TILE_SZ * KV_TILE_SZ, KQ, logit_softcap);
+            }
+
            if (mask) {
                ggml_vec_add_f32(tile_rows * KV_TILE_SZ, KQ, KQ, mask32);
            }
@@ -9549,7 +9320,7 @@ static void ggml_compute_forward_flash_attn_ext_f16(
                                kv_is_f32_or_f16 &&
                                k->type == v->type &&
                                neq1 >= Q_TILE_SZ);
-#if defined(GGML_SIMD) && !defined(__x86_64__) && !defined(_M_X64)
+#ifdef GGML_SIMD
 #if defined(__ARM_FEATURE_SVE)
        const int64_t f32_epr = svcntw();
 #else
@@ -56,56 +56,6 @@ static inline void simd_gemm_ukernel(
    }
 }

-template <int RM>
-static inline void simd_gemm_ukernel_tail(
-    float       * GGML_RESTRICT C,
-    const float * GGML_RESTRICT A,
-    const float * GGML_RESTRICT B,
-    int K, int N, int cols)
-{
-#if defined(__AVX512F__)
-    const __mmask16 mask = (1u << cols) - 1;
-    __m512 acc[RM];
-    for (int64_t i = 0; i < RM; i++) {
-        acc[i] = _mm512_maskz_loadu_ps(mask, C + i * N);
-    }
-    for (int64_t kk = 0; kk < K; kk++) {
-        const __m512 b = _mm512_maskz_loadu_ps(mask, B + kk * N);
-        for (int64_t i = 0; i < RM; i++) {
-            acc[i] = _mm512_mask3_fmadd_ps(_mm512_set1_ps(A[i * K + kk]), b, acc[i], mask);
-        }
-    }
-    for (int64_t i = 0; i < RM; i++) {
-        _mm512_mask_storeu_ps(C + i * N, mask, acc[i]);
-    }
-#elif defined(__AVX2__)
-    const __m256i mask = _mm256_cmpgt_epi32(_mm256_set1_epi32(cols), _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7));
-    __m256 acc[RM];
-    for (int64_t i = 0; i < RM; i++) {
-        acc[i] = _mm256_maskload_ps(C + i * N, mask);
-    }
-    for (int64_t kk = 0; kk < K; kk++) {
-        const __m256 b = _mm256_maskload_ps(B + kk * N, mask);
-        for (int64_t i = 0; i < RM; i++) {
-            acc[i] = GGML_F32_VEC_FMA(acc[i], b, _mm256_set1_ps(A[i * K + kk]));
-        }
-    }
-    for (int64_t i = 0; i < RM; i++) {
-        _mm256_maskstore_ps(C + i * N, mask, acc[i]);
-    }
-#else
-    for (int64_t j = 0; j < cols; j++) {
-        for (int64_t i = 0; i < RM; i++) {
-            float a = C[i * N + j];
-            for (int64_t kk = 0; kk < K; kk++) {
-                a += A[i * K + kk] * B[kk * N + j];
-            }
-            C[i * N + j] = a;
-        }
-    }
-#endif
-}
-
 // C[M x N] += A[M x K] * B[K x N]
 static void simd_gemm(
    float       * GGML_RESTRICT C,
@@ -124,8 +74,14 @@ static void simd_gemm(
        for (; jj + KN <= N; jj += KN) {
            simd_gemm_ukernel<GEMM_RM, 1>(C + jj, A, B + jj, K, N);
        }
-        if (jj < N) {
-            simd_gemm_ukernel_tail<GEMM_RM>(C + jj, A, B + jj, K, N, N - jj);
+        for (; jj < N; jj++) {
+            for (int64_t i = 0; i < GEMM_RM; i++) {
+                float a = C[i * N + jj];
+                for (int64_t kk = 0; kk < K; kk++) {
+                    a += A[i * K + kk] * B[kk * N + jj];
+                }
+                C[i * N + jj] = a;
+            }
        }

        A += GEMM_RM * K;
@@ -141,8 +97,12 @@ static void simd_gemm(
        for (; jj + KN <= N; jj += KN) {
            simd_gemm_ukernel<1, 1>(C + jj, A, B + jj, K, N);
        }
-        if (jj < N) {
-            simd_gemm_ukernel_tail<1>(C + jj, A, B + jj, K, N, N - jj);
+        for (; jj < N; jj++) {
+            float a = C[jj];
+            for (int64_t kk = 0; kk < K; kk++) {
+                a += A[kk] * B[kk * N + jj];
+            }
+            C[jj] = a;
        }

        A += K;
@@ -1314,33 +1314,6 @@ static inline void __lzs_f16cx4_store(ggml_fp16_t * x, float32x4_t v_y) {
 #define GGML_F16_ARR (GGML_F16_STEP/GGML_F16_EPR)
 #endif

-// GGML_F16_DOT_*
-// like GGML_F16_* but for dot products which need F32 accumulation on AVX512-FP16
-
-#if defined(__AVX512FP16__)
-
-#define GGML_F16_DOT_STEP           GGML_F32_STEP
-#define GGML_F16_DOT_EPR            GGML_F32_EPR
-#define GGML_F16_DOT_ARR            GGML_F32_ARR
-#define GGML_F16_DOT_VEC            GGML_F32x16
-#define GGML_F16_DOT_VEC_ZERO       GGML_F32x16_ZERO
-#define GGML_F16_DOT_VEC_LOAD(p, i) _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(p)))
-#define GGML_F16_DOT_VEC_FMA        GGML_F32x16_FMA
-#define GGML_F16_DOT_VEC_REDUCE     GGML_F32x16_REDUCE
-
-#else
-
-#define GGML_F16_DOT_STEP           GGML_F16_STEP
-#define GGML_F16_DOT_EPR            GGML_F16_EPR
-#define GGML_F16_DOT_ARR            GGML_F16_ARR
-#define GGML_F16_DOT_VEC            GGML_F16_VEC
-#define GGML_F16_DOT_VEC_ZERO       GGML_F16_VEC_ZERO
-#define GGML_F16_DOT_VEC_LOAD       GGML_F16_VEC_LOAD
-#define GGML_F16_DOT_VEC_FMA        GGML_F16_VEC_FMA
-#define GGML_F16_DOT_VEC_REDUCE     GGML_F16_VEC_REDUCE
-
-#endif // defined(__AVX512FP16__)
-
 #ifdef __cplusplus
 }
 #endif
@@ -9,7 +9,7 @@
 // bit unpacking routines: tiled_unpk_nib4, tiled_unpk_2bit, tiled_unpk_or
 // LUT value expansion routines: tiled_lut8, tiled_unpk_sign32, tiled_unpk_tern8

-#define GGML_COMMON_DECL_CPP
+#define GGML_COMMON_DECL_C
 #include "ggml-common.h"

 #include <stddef.h>
@@ -35,8 +35,8 @@ static void tiled_unpack_src0(const block_q4_K * rows, int64_t row_stride, int n
            const int s_off = r * nb_stride + slab * NB;
            const block_q4_K & x = rows[r * row_stride + slab];

-            tile->d[d_off]    = ggml_fp16_to_fp32(x.data.data.d);
-            tile->dmin[d_off] = ggml_fp16_to_fp32(x.data.data.dmin);
+            tile->d[d_off]    = ggml_fp16_to_fp32(x.d);
+            tile->dmin[d_off] = ggml_fp16_to_fp32(x.dmin);

            uint32_t utmp[4];
            memcpy(utmp, x.scales, 12);
@@ -80,8 +80,8 @@ static void tiled_unpack_src0(const block_q5_K * rows, int64_t row_stride, int n
            const int s_off = r * nb_stride + slab * NB;
            const block_q5_K & x = rows[r * row_stride + slab];

-            tile->d[d_off]    = ggml_fp16_to_fp32(x.data.data.d);
-            tile->dmin[d_off] = ggml_fp16_to_fp32(x.data.data.dmin);
+            tile->d[d_off]    = ggml_fp16_to_fp32(x.d);
+            tile->dmin[d_off] = ggml_fp16_to_fp32(x.dmin);

            uint32_t utmp[4];
            memcpy(utmp, x.scales, 12);
@@ -207,8 +207,8 @@ static void tiled_unpack_src0(const block_q2_K * rows, int64_t row_stride, int n
            const int q_off = r * qk_stride + slab * TILED_TILE_K;
            const int s_off = r * nb_stride + slab * NB;
            const block_q2_K & x = rows[r * row_stride + slab];
-            tile->d[d_off]    = ggml_fp16_to_fp32(x.data.data.d);
-            tile->dmin[d_off] = ggml_fp16_to_fp32(x.data.data.dmin);
+            tile->d[d_off]    = ggml_fp16_to_fp32(x.d);
+            tile->dmin[d_off] = ggml_fp16_to_fp32(x.dmin);

            // 2-bit code: element e -> half=e>>7, el=e&127
            //   byte = half*32 + (el & 31), shift = 2*(el >> 5)
@@ -342,24 +342,24 @@ void ggml_vec_dot_f16(int n, float * GGML_RESTRICT s, size_t bs, ggml_fp16_t * G
            }
        #endif // __riscv_zvfh
    #else
-        const int np = (n & ~(GGML_F16_DOT_STEP - 1));
+        const int np = (n & ~(GGML_F16_STEP - 1));

-        GGML_F16_DOT_VEC sum[GGML_F16_DOT_ARR] = { GGML_F16_DOT_VEC_ZERO };
+        GGML_F16_VEC sum[GGML_F16_ARR] = { GGML_F16_VEC_ZERO };

-        GGML_F16_DOT_VEC ax[GGML_F16_DOT_ARR];
-        GGML_F16_DOT_VEC ay[GGML_F16_DOT_ARR];
+        GGML_F16_VEC ax[GGML_F16_ARR];
+        GGML_F16_VEC ay[GGML_F16_ARR];

-        for (int i = 0; i < np; i += GGML_F16_DOT_STEP) {
-            for (int j = 0; j < GGML_F16_DOT_ARR; j++) {
-                ax[j] = GGML_F16_DOT_VEC_LOAD(x + i + j*GGML_F16_DOT_EPR, j);
-                ay[j] = GGML_F16_DOT_VEC_LOAD(y + i + j*GGML_F16_DOT_EPR, j);
+        for (int i = 0; i < np; i += GGML_F16_STEP) {
+            for (int j = 0; j < GGML_F16_ARR; j++) {
+                ax[j] = GGML_F16_VEC_LOAD(x + i + j*GGML_F16_EPR, j);
+                ay[j] = GGML_F16_VEC_LOAD(y + i + j*GGML_F16_EPR, j);

-                sum[j] = GGML_F16_DOT_VEC_FMA(sum[j], ax[j], ay[j]);
+                sum[j] = GGML_F16_VEC_FMA(sum[j], ax[j], ay[j]);
            }
        }

        // reduce sum0..sum3 to sum0
-        GGML_F16_DOT_VEC_REDUCE(sumf, sum);
+        GGML_F16_VEC_REDUCE(sumf, sum);

        // leftovers
        for (int i = np; i < n; ++i) {
@@ -276,28 +276,28 @@ inline static void ggml_vec_dot_f16_unroll(const int n, const int xs, float * GG
            const int np = 0;
        #endif
    #else
-        const int np = (n & ~(GGML_F16_DOT_STEP - 1));
+        const int np = (n & ~(GGML_F16_STEP - 1));

-        GGML_F16_DOT_VEC sum[GGML_VEC_DOT_UNROLL][GGML_F16_DOT_ARR] = { { GGML_F16_DOT_VEC_ZERO } };
+        GGML_F16_VEC sum[GGML_VEC_DOT_UNROLL][GGML_F16_ARR] = { { GGML_F16_VEC_ZERO } };

-        GGML_F16_DOT_VEC ax[GGML_F16_DOT_ARR];
-        GGML_F16_DOT_VEC ay[GGML_F16_DOT_ARR];
+        GGML_F16_VEC ax[GGML_F16_ARR];
+        GGML_F16_VEC ay[GGML_F16_ARR];

-        for (int i = 0; i < np; i += GGML_F16_DOT_STEP) {
-            for (int j = 0; j < GGML_F16_DOT_ARR; j++) {
-                ay[j] = GGML_F16_DOT_VEC_LOAD(y + i + j*GGML_F16_DOT_EPR, j);
+        for (int i = 0; i < np; i += GGML_F16_STEP) {
+            for (int j = 0; j < GGML_F16_ARR; j++) {
+                ay[j] = GGML_F16_VEC_LOAD(y + i + j*GGML_F16_EPR, j);

                for (int k = 0; k < GGML_VEC_DOT_UNROLL; ++k) {
-                    ax[j] = GGML_F16_DOT_VEC_LOAD(x[k] + i + j*GGML_F16_DOT_EPR, j);
+                    ax[j] = GGML_F16_VEC_LOAD(x[k] + i + j*GGML_F16_EPR, j);

-                    sum[k][j] = GGML_F16_DOT_VEC_FMA(sum[k][j], ax[j], ay[j]);
+                    sum[k][j] = GGML_F16_VEC_FMA(sum[k][j], ax[j], ay[j]);
                }
            }
        }

        // reduce sum0..sum3 to sum0
        for (int k = 0; k < GGML_VEC_DOT_UNROLL; ++k) {
-            GGML_F16_DOT_VEC_REDUCE(sumf[k], sum[k]);
+            GGML_F16_VEC_REDUCE(sumf[k], sum[k]);
        }
    #endif
 #else
@@ -976,12 +976,6 @@ inline static void ggml_vec_gelu_f16(const int n, ggml_fp16_t * y, const ggml_fp
    }
 }

-inline static void ggml_vec_gelu_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x) {
-    for (int i = 0; i < n; ++i) {
-        y[i] = GGML_FP32_TO_BF16(ggml_gelu_f32(GGML_BF16_TO_FP32(x[i])));
-    }
-}
-
 inline static void ggml_vec_gelu_erf_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
    for (int i = 0; i < n; ++i) {
        float xi = GGML_CPU_FP16_TO_FP32(x[i]);
@@ -990,13 +984,6 @@ inline static void ggml_vec_gelu_erf_f16(const int n, ggml_fp16_t * y, const ggm
    }
 }

-inline static void ggml_vec_gelu_erf_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x) {
-    for (int i = 0; i < n; ++i) {
-        float xi = GGML_BF16_TO_FP32(x[i]);
-        y[i] = GGML_FP32_TO_BF16(0.5f*xi*(1.0f + erff(xi*SQRT_2_INV)));
-    }
-}
-
 #ifdef GGML_GELU_FP16
 inline static void ggml_vec_gelu_f32(const int n, float * y, const float * x) {
    uint16_t t;
@@ -1038,12 +1025,6 @@ inline static void ggml_vec_gelu_quick_f16(const int n, ggml_fp16_t * y, const g
    }
 }

-inline static void ggml_vec_gelu_quick_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x) {
-    for (int i = 0; i < n; ++i) {
-        y[i] = GGML_FP32_TO_BF16(ggml_gelu_quick_f32(GGML_BF16_TO_FP32(x[i])));
-    }
-}
-
 #ifdef GGML_GELU_QUICK_FP16
 inline static void ggml_vec_gelu_quick_f32(const int n, float * y, const float * x) {
    uint16_t t;
@@ -1394,12 +1375,6 @@ inline static void ggml_vec_silu_f16(const int n, ggml_fp16_t * y, const ggml_fp
    }
 }

-inline static void ggml_vec_silu_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x) {
-    for (int i = 0; i < n; ++i) {
-        y[i] = GGML_FP32_TO_BF16(ggml_silu_f32(GGML_BF16_TO_FP32(x[i])));
-    }
-}
-
 inline static float ggml_silu_backward_f32(float x, float dy) {
    const float s = 1.0f/(1.0f + expf(-x));
    return dy*s*(1.0f + x*(1.0f - s));
@@ -1436,13 +1411,6 @@ inline static void ggml_vec_reglu_f16 (const int n, ggml_fp16_t * y, const ggml_
    }
 }

-inline static void ggml_vec_reglu_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x, const ggml_bf16_t * g) {
-    for (int i = 0; i < n; ++i) {
-        float v = GGML_BF16_TO_FP32(x[i]);
-        y[i] = GGML_FP32_TO_BF16((v > 0.f) ? v * GGML_BF16_TO_FP32(g[i]) : 0.f);
-    }
-}
-
 #ifdef GGML_GELU_FP16
 inline static void ggml_vec_geglu_f32(const int n, float * y, const float * x, const float * g) {
    uint16_t t;
@@ -1474,12 +1442,6 @@ inline static void ggml_vec_geglu_f16(const int n, ggml_fp16_t * y, const ggml_f
    }
 }

-inline static void ggml_vec_geglu_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x, const ggml_bf16_t * g) {
-    for (int i = 0; i < n; ++i) {
-        y[i] = GGML_FP32_TO_BF16(ggml_gelu_f32(GGML_BF16_TO_FP32(x[i])) * GGML_BF16_TO_FP32(g[i]));
-    }
-}
-
 void ggml_vec_swiglu_f32(const int n, float * y, const float * x, const float * g);

 inline static void ggml_vec_swiglu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x, const ggml_fp16_t * g) {
@@ -1490,14 +1452,6 @@ inline static void ggml_vec_swiglu_f16(const int n, ggml_fp16_t * y, const ggml_
    }
 }

-inline static void ggml_vec_swiglu_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x, const ggml_bf16_t * g) {
-    for (int i = 0; i < n; ++i) {
-        float xi = GGML_BF16_TO_FP32(x[i]);
-        float gi = GGML_BF16_TO_FP32(g[i]);
-        y[i] = GGML_FP32_TO_BF16((xi/(1.0f + expf(-xi))) * gi);
-    }
-}
-
 inline static void ggml_vec_geglu_erf_f32(const int n, float * y, const float * x, const float * g) {
    for (int i = 0; i < n; ++i) {
        float xi = x[i];
@@ -1513,14 +1467,6 @@ inline static void ggml_vec_geglu_erf_f16(const int n, ggml_fp16_t * y, const gg
    }
 }

-inline static void ggml_vec_geglu_erf_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x, const ggml_bf16_t * g) {
-    for (int i = 0; i < n; ++i) {
-        float xi = GGML_BF16_TO_FP32(x[i]);
-        float gi = GGML_BF16_TO_FP32(g[i]);
-        y[i] = GGML_FP32_TO_BF16(0.5f * xi * (1.0f + erff(xi*SQRT_2_INV)) * gi);
-    }
-}
-
 #ifdef GGML_GELU_QUICK_FP16
 inline static void ggml_vec_geglu_quick_f32(const int n, float * y, const float * x, const float * g) {
    uint16_t t;
@@ -1546,12 +1492,6 @@ inline static void ggml_vec_geglu_quick_f16(const int n, ggml_fp16_t * y, const
    }
 }

-inline static void ggml_vec_geglu_quick_bf16(const int n, ggml_bf16_t * y, const ggml_bf16_t * x, const ggml_bf16_t * g) {
-    for (int i = 0; i < n; ++i) {
-        y[i] = GGML_FP32_TO_BF16(ggml_gelu_quick_f32(GGML_BF16_TO_FP32(x[i])) * GGML_BF16_TO_FP32(g[i]));
-    }
-}
-
 inline static void ggml_vec_sum_f32(const int n, float * s, const float * x) {
 #ifndef GGML_USE_ACCELERATE
    ggml_float sum = 0.0;
@@ -71,13 +71,14 @@ if (CUDAToolkit_FOUND)

    enable_language(CUDA)

-    if (GGML_CUDA_CCCL_VERSION)
+    # TODO: Remove once CCCL 3.2 has been released and bundled with CUDA Toolkit
+    if (GGML_CUDA_CUB_3DOT2)
        include(FetchContent)

        FetchContent_Declare(
            CCCL
            GIT_REPOSITORY https://github.com/nvidia/cccl.git
-            GIT_TAG        "${GGML_CUDA_CCCL_VERSION}"
+            GIT_TAG        v3.2.0
            GIT_SHALLOW    TRUE
        )

@@ -156,15 +157,14 @@ if (CUDAToolkit_FOUND)
        add_compile_definitions(GGML_CUDA_NO_PEER_COPY)
    endif()

-    if (GGML_CUDA_CCCL_VERSION)
-        target_link_libraries(ggml-cuda PRIVATE CCCL::CCCL)
-    endif()
-
    if (GGML_STATIC)
        if (WIN32)
            # As of 12.3.1 CUDA Toolkit for Windows does not offer a static cublas library
            target_link_libraries(ggml-cuda PRIVATE CUDA::cudart_static CUDA::cublas)
        else ()
+            if (GGML_CUDA_CUB_3DOT2)
+                target_link_libraries(ggml-cuda PRIVATE  CCCL::CCCL)
+            endif()
            if (CUDAToolkit_VERSION VERSION_GREATER_EQUAL "10.1")
                target_link_libraries(ggml-cuda PRIVATE  CUDA::cudart_static CUDA::cublas_static CUDA::cublasLt_static)
            else()
@@ -172,6 +172,9 @@ if (CUDAToolkit_FOUND)
            endif()
        endif()
    else()
+        if (GGML_CUDA_CUB_3DOT2)
+            target_link_libraries(ggml-cuda PRIVATE  CCCL::CCCL)
+        endif()
        target_link_libraries(ggml-cuda PRIVATE CUDA::cudart CUDA::cublas)
    endif()

@@ -166,62 +166,52 @@ static inline __device__ void ggml_cuda_swap(T & a, T & b) {
    b = tmp;
 }

-// One compare-exchange of the bitonic network at (k, j) for column col.
-template<ggml_sort_order order>
-static inline __device__ void bitonic_step(const float * x_row, int * dst_row, const int ncols, const int col, const int k, const int j) {
-    const int ixj = col ^ j;
-    if (ixj <= col) {
-        return;
-    }
-    if ((col & k) == 0) {
-        if (dst_row[col] >= ncols ||
-            (dst_row[ixj] < ncols && (order == GGML_SORT_ORDER_ASC ?
-                x_row[dst_row[col]] > x_row[dst_row[ixj]] :
-                x_row[dst_row[col]] < x_row[dst_row[ixj]]))
-        ) {
-            ggml_cuda_swap(dst_row[col], dst_row[ixj]);
-        }
-    } else {
-        if (dst_row[ixj] >= ncols ||
-            (dst_row[col] < ncols && (order == GGML_SORT_ORDER_ASC ?
-                x_row[dst_row[col]] < x_row[dst_row[ixj]] :
-                x_row[dst_row[col]] > x_row[dst_row[ixj]]))
-        ) {
-            ggml_cuda_swap(dst_row[col], dst_row[ixj]);
-        }
-    }
-}
-
-// Bitonic sort of one row per block. Each thread owns the columns
-// threadIdx.x + i * blockDim.x, so rows wider than the block (up to the
-// shared memory limit) sort with several columns per thread. Every
-// (k, j) stage runs all owned columns before the barrier; a pair
-// (col, col ^ j) is exchanged by the owner of its lower index only.
 template<ggml_sort_order order>
 static __global__ void k_argsort_f32_i32(const float * x, int * dst, const int ncols, int ncols_pad) {
-    const int row = blockIdx.x;
+    // bitonic sort
+    int col = threadIdx.x;
+    int row = blockIdx.x;
+
+    if (col >= ncols_pad) {
+        return;
+    }

    const float * x_row = x + row * ncols;
    extern __shared__ int dst_row[];

    // initialize indices
-    for (int col = threadIdx.x; col < ncols_pad; col += blockDim.x) {
-        dst_row[col] = col;
-    }
+    dst_row[col] = col;

    __syncthreads();

    for (int k = 2; k <= ncols_pad; k *= 2) {
        for (int j = k / 2; j > 0; j /= 2) {
-            for (int col = threadIdx.x; col < ncols_pad; col += blockDim.x) {
-                bitonic_step<order>(x_row, dst_row, ncols, col, k, j);
+            int ixj = col ^ j;
+            if (ixj > col) {
+                if ((col & k) == 0) {
+                    if (dst_row[col] >= ncols ||
+                        (dst_row[ixj] < ncols && (order == GGML_SORT_ORDER_ASC ?
+                            x_row[dst_row[col]] > x_row[dst_row[ixj]] :
+                            x_row[dst_row[col]] < x_row[dst_row[ixj]]))
+                    ) {
+                        ggml_cuda_swap(dst_row[col], dst_row[ixj]);
+                    }
+                } else {
+                    if (dst_row[ixj] >= ncols ||
+                        (dst_row[col] < ncols && (order == GGML_SORT_ORDER_ASC ?
+                            x_row[dst_row[col]] < x_row[dst_row[ixj]] :
+                            x_row[dst_row[col]] > x_row[dst_row[ixj]]))
+                    ) {
+                        ggml_cuda_swap(dst_row[col], dst_row[ixj]);
+                    }
+                }
            }
            __syncthreads();
        }
    }

    // copy the result to dst without the padding
-    for (int col = threadIdx.x; col < ncols; col += blockDim.x) {
+    if (col < ncols) {
        dst[row * ncols + col] = dst_row[col];
    }
 }
@@ -243,9 +233,7 @@ void argsort_f32_i32_cuda_bitonic(const float *   x,
    // bitonic sort requires ncols to be power of 2
    const int ncols_pad = next_power_of_2(ncols);

-    // one thread per column up to the block limit, several columns per
-    // thread beyond it; shared memory is the remaining bound
-    const dim3 block_dims(ncols_pad < CUDA_ARGSORT_BLOCK_SIZE ? ncols_pad : CUDA_ARGSORT_BLOCK_SIZE, 1, 1);
+    const dim3 block_dims(ncols_pad, 1, 1);
    const dim3 block_nums(nrows, 1, 1);
    const size_t shared_mem = ncols_pad * sizeof(int);

@@ -1,7 +1,5 @@
 #include "common.cuh"

-#define CUDA_ARGSORT_BLOCK_SIZE 1024
-
 void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst);

 #ifdef GGML_CUDA_USE_CUB
@@ -1,5 +1,4 @@
 #include "binbcast.cuh"
-#include "convert.cuh"
 #include <cstdint>
 #include <utility>

@@ -83,14 +82,14 @@ static __global__ void k_bin_bcast(const src0_t *         src0,
    for (uint32_t i0 = i0s; i0 < ne0; i0 += s0) {
        const uint32_t i10 = fastmodulo(i0, ne10);

-        float result = src0_row ? ggml_cuda_cast<float>(src0_row[size_t(i0)*s00]) : 0.0f;
+        float result = src0_row ? (float) src0_row[size_t(i0)*s00] : 0.0f;
        if constexpr (sizeof...(src1_ptrs) > 0) {
-            result = (..., (result = bin_op(result, ggml_cuda_cast<float>(src1s[i_src1 + size_t(i10)*s10]))));
+            result = (..., (result = bin_op(result, (float)src1s[i_src1 + size_t(i10)*s10])));
        } else {
-            result = bin_op(result, ggml_cuda_cast<float>(src1[i_src1 + size_t(i10)*s10]));
+            result = bin_op(result, (float)src1[i_src1 + size_t(i10)*s10]);
        }

-        dst_row[i0] = ggml_cuda_cast<dst_t>(result);
+        dst_row[i0] = (dst_t) result;

        // protect i0 from overflow
        if (ne0 - i0 <= s0) {
@@ -155,14 +154,14 @@ static __global__ void k_bin_bcast_unravel(const src0_t *         src0,
    const uint32_t i10 = fastmodulo(i0, ne10);

    ggml_cuda_pdl_sync();
-    float result = src0_row ? ggml_cuda_cast<float>(src0_row[size_t(i0)*s00]) : 0.0f;
+    float result = src0_row ? (float) src0_row[size_t(i0)*s00] : 0.0f;
    if constexpr (sizeof...(src1_ptrs) > 0) {
-        result = (..., (result = bin_op(result, ggml_cuda_cast<float>(src1s[i_src1 + size_t(i10)*s10]))));
+        result = (..., (result = bin_op(result, (float)src1s[i_src1 + size_t(i10)*s10])));
    } else {
-        result = bin_op(result, ggml_cuda_cast<float>(src1[i_src1 + size_t(i10)*s10]));
+        result = bin_op(result, (float)src1[i_src1 + size_t(i10)*s10]);
    }

-    dst_row[i0] = ggml_cuda_cast<dst_t>(result);
+    dst_row[i0] = (dst_t) result;
 }

 template <float (*bin_op)(const float, const float), typename src0_t, typename src1_t, typename dst_t, size_t... I>
@@ -414,14 +413,10 @@ static void ggml_cuda_op_bin_bcast(
    const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst,
    const void * src0_dd, const void * src1_dd, void * dst_dd, cudaStream_t stream) {

-    GGML_ASSERT(src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16 || src1->type == GGML_TYPE_BF16);
+    GGML_ASSERT(src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16);

-    if (src0->type == GGML_TYPE_F32 && src1->type != GGML_TYPE_BF16 && dst->type == GGML_TYPE_F32) {
+    if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
        op()(src0, src1, dst, (const float *)src0_dd, (const float *)src1_dd, (float *)dst_dd, stream);
-    } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_BF16 && dst->type == GGML_TYPE_BF16) {
-        op()(src0, src1, dst, (const nv_bfloat16 *) src0_dd, (const nv_bfloat16 *) src1_dd, (nv_bfloat16 *) dst_dd, stream);
-    } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_BF16) {
-        op()(src0, src1, dst, (const nv_bfloat16 *) src0_dd, (const float *) src1_dd, (nv_bfloat16 *) dst_dd, stream);
    } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) {
        op()(src0, src1, dst, (const half *) src0_dd, (const half *)src1_dd, (half *) dst_dd, stream);
    } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F16) {
@@ -462,18 +457,10 @@ static void ggml_cuda_op_fused_binbcast_impl(ggml_backend_cuda_context & ctx, gg
    const ggml_tensor * src0 = dst->src[0];
    const ggml_tensor * src1 = dst->src[1];

-    if (src0->type == GGML_TYPE_F32 && src1->type != GGML_TYPE_BF16 && dst->type == GGML_TYPE_F32) {
+    if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
        launch_bin_bcast_pack<op, float, float, float>(src0, src1, dst,
            (const float *) src0->data, (const float *) src1->data, (float *) dst->data,
            stream, std::make_index_sequence<n_fuse>{});
-    } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_BF16 && dst->type == GGML_TYPE_BF16) {
-        launch_bin_bcast_pack<op, nv_bfloat16, nv_bfloat16, nv_bfloat16>(src0, src1, dst,
-            (const nv_bfloat16 *) src0->data, (const nv_bfloat16 *) src1->data, (nv_bfloat16 *) dst->data,
-            stream, std::make_index_sequence<n_fuse>{});
-    } else if (src0->type == GGML_TYPE_BF16 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_BF16) {
-        launch_bin_bcast_pack<op, nv_bfloat16, float, nv_bfloat16>(src0, src1, dst,
-            (const nv_bfloat16 *) src0->data, (const float *) src1->data, (nv_bfloat16 *) dst->data,
-            stream, std::make_index_sequence<n_fuse>{});
    } else if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) {
        launch_bin_bcast_pack<op, half, half, half>(src0, src1, dst,
            (const half *) src0->data, (const half *) src1->data, (half *) dst->data,
@@ -131,15 +131,15 @@ static __device__ __forceinline__ void ggml_cuda_syncwarp() {
 }

 static __device__ __forceinline__ void ggml_cuda_pdl_sync() {
-#if defined(GGML_CUDA_USE_PDL) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER
+#if defined(GGML_CUDA_USE_PDL) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER
    cudaGridDependencySynchronize();
-#endif // defined(GGML_CUDA_USE_PDL) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER
+#endif // defined(GGML_CUDA_USE_PDL) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER
 }

 static __device__ __forceinline__ void ggml_cuda_pdl_lc() {
-#if defined(GGML_CUDA_USE_PDL) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER
+#if defined(GGML_CUDA_USE_PDL) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER
    cudaTriggerProgrammaticLaunchCompletion();
-#endif // defined(GGML_CUDA_USE_PDL) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER
+#endif // defined(GGML_CUDA_USE_PDL) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER
 }

 #ifdef __CUDA_ARCH_LIST__
@@ -285,21 +285,21 @@ static const char * cu_get_error_str(CUresult err) {
 #define VOLTA_MMA_AVAILABLE
 #endif // !defined(GGML_USE_HIP) && __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA

-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_TURING
+#if !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_TURING
 #define TURING_MMA_AVAILABLE
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_TURING
+#endif // !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_TURING

-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+#if !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
 #define AMPERE_MMA_AVAILABLE
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+#endif // !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE

 #if !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_BLACKWELL && __CUDA_ARCH__ < GGML_CUDA_CC_RUBIN
 #    define BLACKWELL_MMA_AVAILABLE
 #endif // !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_BLACKWELL

-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+#if !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
 #define CP_ASYNC_AVAILABLE
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+#endif // !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE

 #if !defined(GGML_CUDA_NO_FA) && !(defined(GGML_USE_MUSA) && __MUSA_ARCH__ < 220)
 #define FLASH_ATTN_AVAILABLE
@@ -453,7 +453,7 @@ struct ggml_cuda_unroll<1> {

 template<int width = WARP_SIZE>
 static __device__ __forceinline__ int warp_reduce_sum(int x) {
-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+#if !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
    return __reduce_add_sync(0xffffffff, x);
 #else
 #pragma unroll
@@ -461,7 +461,7 @@ static __device__ __forceinline__ int warp_reduce_sum(int x) {
        x += __shfl_xor_sync(0xffffffff, x, offset, width);
    }
    return x;
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
+#endif // !defined(GGML_USE_HIP) && __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
 }

 template<int width = WARP_SIZE>
@@ -1680,11 +1680,11 @@ static bool ggml_cuda_kernel_can_use_pdl(const void * kernel) {
 #endif //defined(GGML_CUDA_USE_PDL)

 // PDL and __restrict__ need to be mutually exclusive, see https://github.com/ggml-org/llama.cpp/pull/24030
-# if (defined(GGML_CUDA_USE_PDL) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER)
+# if (defined(GGML_CUDA_USE_PDL) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER)
 # define GGML_CUDA_RESTRICT
 # else
 # define GGML_CUDA_RESTRICT __restrict__
-# endif // defined(GGML_CUDA_USE_PDL) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER
+# endif // defined(GGML_CUDA_USE_PDL) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER

 template<typename Kernel, typename... Args>
 static __inline__ void ggml_cuda_kernel_launch(Kernel kernel, const ggml_cuda_kernel_launch_params & launch_params, Args&&... args) {
@@ -223,14 +223,9 @@ static __global__ void dequantize_block_iq1_m(const void * __restrict__ vx, dst_
 }

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

-    const int64_t ib = 8*i + threadIdx.x%8;
-    if (ib >= nb32) {
-        return;
-    }
-
    dequantize_iq4_nl(vx, i, yy + i*QK_K, threadIdx.x);
 }

@@ -357,9 +352,8 @@ static void dequantize_row_iq1_s_cuda(const void * vx, dst_t * y, const int64_t

 template<typename dst_t>
 static void dequantize_row_iq4_nl_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) {
-    const int nb32 = k / QK4_NL;
    const int nb = (k + QK_K - 1) / QK_K;
-    dequantize_block_iq4_nl<<<nb, 32, 0, stream>>>(vx, y, nb32);
+    dequantize_block_iq4_nl<<<nb, 32, 0, stream>>>(vx, y);
 }

 template<typename dst_t>
@@ -45,11 +45,11 @@ template<typename dst_t, typename src_t>
 #ifdef GGML_USE_HIP
        return make_float2(__bfloat162float(__low2bfloat16(x)), __bfloat162float(__high2bfloat16(x)));
 #else
-#if !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= 800
+#if __CUDA_ARCH__ >= 800
        return __bfloat1622float2(x);
 #else
        return make_float2(__bfloat162float(x.x), __bfloat162float(x.y));
-#endif // !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= 800
+#endif // __CUDA_ARCH__ >= 800
 #endif // GGML_USE_HIP
    } else if constexpr(std::is_same_v<src_t, float2> && std::is_same_v<dst_t, nv_bfloat162>) {
        // bypass compile error on cuda 12.0.1
@@ -416,14 +416,9 @@ static __device__ __forceinline__ void dequantize_V_q4_0(const void * __restrict
    int q;
    static_assert(ne == 2 || ne == 4, "bad ne");
    ggml_cuda_memcpy_1<ne, 2>(&q, x[ib].qs + iqs);
-#if defined(GGML_USE_HIP)
-    // Keep this VMEM read close to its packed-byte dequantization. Hoisting it too far
-    // increases VGPR pressure substantially in some FlashAttention vector kernels.
-    __builtin_amdgcn_sched_group_barrier(0x20, 1, 0);
-#endif // defined(GGML_USE_HIP)
    q >>= 4*shift;
    q &= 0x0F0F0F0F;
-    q = __vsub4(q, 0x08080808);
+    q = __vsubss4(q, 0x08080808);

    const int8_t * q8 = (const int8_t *) &q;

@@ -513,7 +508,7 @@ static __device__ __forceinline__ void dequantize_V_q5_0(const void * __restrict
        }
    }

-    q = __vsub4(q, 0x10101010);
+    q = __vsubss4(q, 0x10101010);

    const int8_t * q8 = (const int8_t *) &q;

@@ -365,7 +365,7 @@ static void ggml_cuda_flash_attn_ext_mma_f16(ggml_backend_cuda_context & ctx, gg

            GGML_ASSERT(Q->ne[2] % K->ne[2] == 0);
            const int gqa_ratio = Q->ne[2] / K->ne[2];
-            if (gqa_ratio == 20 && GGML_CUDA_CC_IS_NVIDIA(cc)) { // GLM 4.7 Flash
+            if (gqa_ratio == 20) { // GLM 4.7 Flash
                if (cc >= GGML_CUDA_CC_DGX_SPARK) {
                    if (Q->ne[1] <= 8) {
                        ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<576, 512, 16>(ctx, dst);
@@ -1616,7 +1616,6 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const

 static void ggml_cuda_mul_mat_cublas(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
    const int cc = ggml_cuda_info().devices[ctx.device].cc;
-    const ggml_prec prec = (ggml_prec) ggml_get_op_params_i32(dst, 0);
    ggml_type compute_type = src0->type;
    if (ggml_is_quantized(compute_type)) {
        compute_type = fast_fp16_hardware_available(cc) ? GGML_TYPE_F16 : GGML_TYPE_F32;
@@ -1630,10 +1629,7 @@ static void ggml_cuda_mul_mat_cublas(ggml_backend_cuda_context & ctx, const ggml
            compute_type = GGML_TYPE_F32;
        }
    }
-    // F16 is the only compute type that can not satisfy a request for BF16
-    if (prec == GGML_PREC_BF16 && compute_type == GGML_TYPE_F16) {
-        compute_type = fast_bf16_hardware_available(cc) ? GGML_TYPE_BF16 : GGML_TYPE_F32;
-    } else if (prec == GGML_PREC_F32) {
+    if (dst->op_params[0] == GGML_PREC_F32) {
        compute_type = GGML_TYPE_F32;
    }

@@ -2037,7 +2033,6 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor *

        ggml_tensor dst_slice;
        memset(&dst_slice, 0, sizeof(dst_slice));
-        memcpy(dst_slice.op_params, dst->op_params, sizeof(dst_slice.op_params));
        dst_slice.buffer = dst->buffer;
        dst_slice.type   = type_dst_sorted;
        dst_slice.ne[0]  = ne0;
@@ -3376,7 +3371,7 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph *                cgraph,
            return false;
        }

-        if (unary->type != GGML_TYPE_F32 && unary->type != GGML_TYPE_F16 && unary->type != GGML_TYPE_BF16) {
+        if (unary->type != GGML_TYPE_F32 && unary->type != GGML_TYPE_F16) {
            return false;
        }

@@ -3425,10 +3420,8 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph *                cgraph,
        const ggml_tensor *tanh   = cgraph->nodes[node_idx+1];
        const ggml_tensor *scale2 = cgraph->nodes[node_idx+2];

-        // the fused softcap kernel is F32 only
-        if (scale->src[0]->type != GGML_TYPE_F32 || scale->type != GGML_TYPE_F32) {
-            return false;
-        }
+        GGML_ASSERT(scale->src[0]->type == GGML_TYPE_F32);
+        GGML_ASSERT(scale->type == GGML_TYPE_F32);

        if (ggml_get_unary_op(tanh) != GGML_UNARY_OP_TANH) {
            return false;
@@ -5189,9 +5182,6 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                case GGML_UNARY_OP_CEIL:
                case GGML_UNARY_OP_ROUND:
                case GGML_UNARY_OP_TRUNC:
-                    if (op->src[0]->type == GGML_TYPE_BF16 && ggml_get_unary_op(op) == GGML_UNARY_OP_XIELU) {
-                        return false;
-                    }
                    // TODO: should become:
                    //return ggml_is_contiguous_rows(op->src[0]);
                    return ggml_is_contiguous(op->src[0]);
@@ -5208,10 +5198,6 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                case GGML_GLU_OP_GEGLU_ERF:
                case GGML_GLU_OP_GEGLU_QUICK:
                case GGML_GLU_OP_SWIGLU_CLAMP:
-                    if (op->src[0]->type == GGML_TYPE_BF16 &&
-                            (ggml_get_glu_op(op) == GGML_GLU_OP_SWIGLU_OAI || ggml_get_glu_op(op) == GGML_GLU_OP_SWIGLU_CLAMP)) {
-                        return false;
-                    }
                    return ggml_is_contiguous_1(op->src[0]);
                default:
                    return false;
@@ -5481,6 +5467,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
        case GGML_OP_TRANSPOSE:
        case GGML_OP_ADD_ID:
        case GGML_OP_ADD1:
+        case GGML_OP_SCALE:
        case GGML_OP_SQR:
        case GGML_OP_SQRT:
        case GGML_OP_SIN:
@@ -5488,16 +5475,10 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
        case GGML_OP_CLAMP:
        case GGML_OP_LOG:
            return true;
-        case GGML_OP_SCALE:
-            return (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_BF16) && op->type == op->src[0]->type;
        case GGML_OP_ADD:
        case GGML_OP_SUB:
        case GGML_OP_MUL:
        case GGML_OP_DIV:
-            if (op->src[0]->type == GGML_TYPE_BF16 || op->src[1]->type == GGML_TYPE_BF16 || op->type == GGML_TYPE_BF16) {
-                return op->src[0]->type == GGML_TYPE_BF16 && op->type == GGML_TYPE_BF16 &&
-                    (op->src[1]->type == GGML_TYPE_BF16 || op->src[1]->type == GGML_TYPE_F32);
-            }
            return (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16) &&
                   (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) &&
                   (op->type         == GGML_TYPE_F32 || op->type         == GGML_TYPE_F16);
@@ -5572,14 +5553,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
 #endif // defined(GGML_USE_HIP) || defined(GGML_CUDA_USE_CUB)
        case GGML_OP_ARGSORT:
 #ifndef GGML_CUDA_USE_CUB
-            {
-                // bitonic path: the padded row must fit in shared memory
-                int64_t ncols_pad = 1;
-                while (ncols_pad < op->src[0]->ne[0]) {
-                    ncols_pad *= 2;
-                }
-                return ncols_pad * sizeof(int) <= ggml_cuda_info().devices[dev_ctx->device].smpb;
-            }
+            return op->src[0]->ne[0] <= 1024;
 #else
            return true;
 #endif
@@ -218,8 +218,8 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
        const int qs0 = get_int_b2(bxi->qs, kqsx);

 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-        x_qs[i*sram_stride + kbx*(2*QI4_0) + kqsx + 0]     = __vsub4((qs0 >> 0) & 0x0F0F0F0F, 0x08080808);
-        x_qs[i*sram_stride + kbx*(2*QI4_0) + kqsx + QI4_0] = __vsub4((qs0 >> 4) & 0x0F0F0F0F, 0x08080808);
+        x_qs[i*sram_stride + kbx*(2*QI4_0) + kqsx + 0]     = __vsubss4((qs0 >> 0) & 0x0F0F0F0F, 0x08080808);
+        x_qs[i*sram_stride + kbx*(2*QI4_0) + kqsx + QI4_0] = __vsubss4((qs0 >> 4) & 0x0F0F0F0F, 0x08080808);
 #else
        x_qs[i*(MMQ_TILE_NE_K + 1) + txi] = qs0;
 #endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)
@@ -350,14 +350,14 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
        qs0    |= (qh << 11)   & 0x00001000;  // 1 -> 12
        qs0    |= (qh << 18)   & 0x00100000;  // 2 -> 20
        qs0    |= (qh << 25)   & 0x10000000;  // 3 -> 28
-        qs0     = __vsub4(qs0, 0x10101010); // subtract 16
+        qs0     = __vsubss4(qs0, 0x10101010); // subtract 16

        int qs1 = (ql >>  4)   & 0x0F0F0F0F;
        qs1    |= (qh >> 12)   & 0x00000010;  // 16 ->  4
        qs1    |= (qh >>  5)   & 0x00001000;  // 17 -> 12
        qs1    |= (qh <<  2)   & 0x00100000;  // 18 -> 20
        qs1    |= (qh <<  9)   & 0x10000000;  // 19 -> 28
-        qs1     = __vsub4(qs1, 0x10101010); // subtract 16
+        qs1     = __vsubss4(qs1, 0x10101010); // subtract 16

 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
        x_qs[i*sram_stride + kbx*(2*QI5_0) + kqsx + 0]     = qs0;
@@ -636,7 +636,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
            const int x_ql_k =  (x_ql_0 >> (2*l))       & 0x03030303;
            const int x_qh_k = ((x_qh_0 >>    l)  << 2) & 0x04040404;

-            const int x_qs_k = __vsub4(x_ql_k | x_qh_k, 0x04040404);
+            const int x_qs_k = __vsubss4(x_ql_k | x_qh_k, 0x04040404);

 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
            x_qs[i*sram_stride           + k] = x_qs_k;
@@ -667,7 +667,7 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
        const int shift_high = 2 * ksc;
        const int sc_high = ((get_int_b2(bxi->scales, ksc_high) >> shift_high) << 4) & 0x30303030;

-        const int sc = __vsub4(sc_low | sc_high, 0x20202020);
+        const int sc = __vsubss4(sc_low | sc_high, 0x20202020);

 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
        const int8_t * sc8 = (const int8_t *) &sc;
@@ -987,11 +987,11 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
        const int kq1 = 2*txi - txi % (QI6_K/2) + QI6_K/2;

 #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
-        x_qs[i*sram_stride + kq0] = __vsub4(ql0 | qh0, 0x20202020);
-        x_qs[i*sram_stride + kq1] = __vsub4(ql1 | qh1, 0x20202020);
+        x_qs[i*sram_stride + kq0] = __vsubss4(ql0 | qh0, 0x20202020);
+        x_qs[i*sram_stride + kq1] = __vsubss4(ql1 | qh1, 0x20202020);
 #else
-        x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq0] = __vsub4(ql0 | qh0, 0x20202020);
-        x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq1] = __vsub4(ql1 | qh1, 0x20202020);
+        x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq0] = __vsubss4(ql0 | qh0, 0x20202020);
+        x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq1] = __vsubss4(ql1 | qh1, 0x20202020);
 #endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
    }

@@ -284,9 +284,9 @@ static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_t
        return ggml_cuda_mmq_get_config_ampere(type, J, fallback);
    }
    return ggml_cuda_mmq_get_config_blackwell(type, J, fallback);
-#elif !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
+#elif __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA
    return ggml_cuda_mmq_get_config_ampere(type, J, fallback);
-#elif !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_DP4A
+#elif __CUDA_ARCH__ >= GGML_CUDA_CC_DP4A
    return ggml_cuda_mmq_get_config_pascal_dp4a(type, J, fallback);
 #else
    return ggml_cuda_mmq_get_config_pascal_older(type, J, fallback);
@@ -9,7 +9,7 @@
 // only enabled on DGX Spark, where it is a gain on every type below. On the higher-bandwidth parts the kernel
 // has little exposed latency left to hide and the extra requests cost more than they save.
 // For perf data, see https://github.com/ggml-org/llama.cpp/pull/26705#issuecomment-5569335031
-#if __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK
 // returns true only for those quants that benefit from prefetch and false otherwise
 static constexpr __host__ __device__ bool mmvq_should_prefetch(ggml_type type) {
    switch (type) {
@@ -112,9 +112,9 @@ static constexpr __device__ mmvq_parameter_table_id get_device_table_id() {
    return MMVQ_PARAMETERS_RDNA2;
 #elif defined(GCN) || defined(CDNA)
    return MMVQ_PARAMETERS_GCN;
-#elif __CUDA_ARCH__ >= GGML_CUDA_CC_TURING && __CUDA_ARCH__ < GGML_CUDA_CC_AMPERE
+#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_TURING && __CUDA_ARCH__ < GGML_CUDA_CC_AMPERE
    return MMVQ_PARAMETERS_TURING;
-#elif __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK
+#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK
    return MMVQ_PARAMETERS_GB10;
 #else
    return MMVQ_PARAMETERS_GENERIC;
@@ -440,9 +440,9 @@ static constexpr __device__ int get_mmvq_mmid_max_batch_for_device() {
    return get_mmvq_mmid_max_batch_cdna(type);
 #elif defined(GCN)
    return get_mmvq_mmid_max_batch_gcn(type);
-#elif !defined(GGML_USE_MUSA) && (__CUDA_ARCH__ == GGML_CUDA_CC_VOLTA || __CUDA_ARCH__ >= GGML_CUDA_CC_ADA_LOVELACE)
+#elif defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == GGML_CUDA_CC_VOLTA || __CUDA_ARCH__ >= GGML_CUDA_CC_ADA_LOVELACE)
    return MMVQ_MAX_BATCH_SIZE;
-#elif !defined(GGML_USE_MUSA) && __CUDA_ARCH__ >= GGML_CUDA_CC_TURING
+#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_TURING
    return get_mmvq_mmid_max_batch_turing_plus(type);
 #else
    return get_mmvq_mmid_max_batch_pascal_older(type);
@@ -718,7 +718,7 @@ static __global__ void mul_mat_vec_q(
        // x block quant index when casting the quants to int
        const int kqs = vdr * (tid % (qi/vdr));

-#if __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK
        // start the next iterations' weight loads early
        if constexpr (mmvq_should_prefetch(type)) {
            constexpr int pf_dist = 2; // loop iterations, not blocks
@@ -1,42 +1,37 @@
 #include "scale.cuh"
-#include "convert.cuh"

 #define MAX_GRIDDIM_X 0x7FFFFFFF

-template <typename T>
-static __global__ void scale_f32(const T * x, T * dst, const float scale, const float bias, const int64_t nelements) {
+static __global__ void scale_f32(const float * x, float * dst, const float scale, const float bias, const int64_t nelements) {
    ggml_cuda_pdl_lc();
    int64_t tid = (int64_t)blockIdx.x * (int64_t)blockDim.x + (int64_t)threadIdx.x;
    int64_t stride = (int64_t)blockDim.x * (int64_t)gridDim.x;

    ggml_cuda_pdl_sync();
    for (int64_t i = tid; i < nelements; i += stride) {
-        dst[i] = ggml_cuda_cast<T>(scale * ggml_cuda_cast<float>(x[i]) + bias);
+        dst[i] = scale * x[i] + bias;
    }
 }

-template <typename T>
-static void scale_f32_cuda(const T * x, T * dst, const float scale, const float bias, const int64_t nelements, cudaStream_t stream) {
+static void scale_f32_cuda(const float * x, float * dst, const float scale, const float bias, const int64_t nelements, cudaStream_t stream) {
    const int64_t num_blocks = (nelements + CUDA_SCALE_BLOCK_SIZE - 1) / CUDA_SCALE_BLOCK_SIZE;
    const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(MIN(MAX_GRIDDIM_X, num_blocks), CUDA_SCALE_BLOCK_SIZE, 0, stream);
-    ggml_cuda_kernel_launch(scale_f32<T>, launch_params, x, dst, scale, bias, nelements);
+    ggml_cuda_kernel_launch(scale_f32, launch_params, x, dst, scale, bias, nelements);
 }

 void ggml_cuda_op_scale(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
    const ggml_tensor * src0 = dst->src[0];
+    const float * src0_d = (const float *)src0->data;
+    float * dst_d = (float *)dst->data;
    cudaStream_t stream = ctx.stream();

-    GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_BF16);
-    GGML_ASSERT(dst->type == src0->type);
+    GGML_ASSERT(src0->type == GGML_TYPE_F32);
+    GGML_ASSERT( dst->type == GGML_TYPE_F32);

    float scale;
    float bias;
    memcpy(&scale, (float *) dst->op_params + 0, sizeof(float));
    memcpy(&bias,  (float *) dst->op_params + 1, sizeof(float));

-    if (src0->type == GGML_TYPE_BF16) {
-        scale_f32_cuda((const nv_bfloat16 *) src0->data, (nv_bfloat16 *) dst->data, scale, bias, ggml_nelements(src0), stream);
-    } else {
-        scale_f32_cuda((const float *) src0->data, (float *) dst->data, scale, bias, ggml_nelements(src0), stream);
-    }
+    scale_f32_cuda(src0_d, dst_d, scale, bias, ggml_nelements(src0), stream);
 }
@@ -3,15 +3,11 @@

 #ifdef GGML_CUDA_USE_CUB
 #    include <cub/cub.cuh>
-// DeviceTopK has a race condition before CCCL 3.4.3.
-// https://github.com/NVIDIA/cccl/pull/10627
-#    if (CCCL_MAJOR_VERSION > 3 || \
-         (CCCL_MAJOR_VERSION == 3 && CCCL_MINOR_VERSION > 4) || \
-         (CCCL_MAJOR_VERSION == 3 && CCCL_MINOR_VERSION == 4 && CCCL_PATCH_VERSION >= 3))
+#    if (CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 2)
 #        define CUB_TOP_K_AVAILABLE
 #        include <cuda/iterator>
 using namespace cub;
-#    endif  // CCCL >= 3.4.3
+#    endif  // CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 2
 #endif      // GGML_CUDA_USE_CUB

 #ifdef CUB_TOP_K_AVAILABLE
@@ -124,7 +124,7 @@ static __global__ void unary_op_kernel(const T * x, T * dst, const int k) {
    }

    ggml_cuda_pdl_sync();
-    dst[i] = ggml_cuda_cast<T>(op(ggml_cuda_cast<float>(x[i])));
+    dst[i] = (T)op((float)x[i]);
 }

 template <float (*op)(float), typename T>
@@ -143,13 +143,12 @@ void ggml_cuda_op_unary(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {

    GGML_ASSERT(ggml_is_contiguous(src0));

-    GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16);
+    GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16);
+    GGML_ASSERT( dst->type == GGML_TYPE_F32 ||  dst->type == GGML_TYPE_F16);
    GGML_ASSERT(src0->type == dst->type);

    if (src0->type == GGML_TYPE_F16) {
        unary_cuda<op>((const half *)src0_d, (half *)dst_d, ggml_nelements(src0), stream);
-    } else if (src0->type == GGML_TYPE_BF16) {
-        unary_cuda<op>((const nv_bfloat16 *)src0_d, (nv_bfloat16 *)dst_d, ggml_nelements(src0), stream);
    } else {
        unary_cuda<op>((const float *)src0_d, (float *)dst_d, ggml_nelements(src0), stream);
    }
@@ -274,7 +273,7 @@ static __global__ void unary_gated_op_kernel(const T * x, const T * g, T * dst,
    const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n);

    ggml_cuda_pdl_sync();
-    dst[i] = ggml_cuda_cast<T>(op(ggml_cuda_cast<float>(x[j0])) * ggml_cuda_cast<float>(g[j1]));
+    dst[i] = (T)(op((float)x[j0]) * (float)g[j1]);
 }

 template <float (*op)(float), typename T>
@@ -300,7 +299,8 @@ void ggml_cuda_op_unary_gated(ggml_backend_cuda_context & ctx, ggml_tensor * dst
    GGML_ASSERT(src0->nb[0] == ggml_element_size(src0));
    GGML_ASSERT(ggml_is_contiguous(dst));

-    GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16);
+    GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16);
+    GGML_ASSERT( dst->type == GGML_TYPE_F32 ||  dst->type == GGML_TYPE_F16);
    GGML_ASSERT(src0->type == dst->type);
    GGML_ASSERT(dst->ne[0] == nc);
    GGML_ASSERT(ggml_nrows(dst) == ggml_nrows(src0));
@@ -324,16 +324,6 @@ void ggml_cuda_op_unary_gated(ggml_backend_cuda_context & ctx, ggml_tensor * dst
        }

        unary_gated_cuda<op>(src0_p, src1_p, (half *)dst_d, ggml_nelements(dst), nc, src0_o / sizeof(half), src1_o / sizeof(half), stream);
-    } else if (src0->type == GGML_TYPE_BF16) {
-        nv_bfloat16 * src0_p = (nv_bfloat16 *) src0_d;
-        nv_bfloat16 * src1_p = (nv_bfloat16 *) src1_d;
-
-        if (!src1) {
-            src0_p += swapped ? nc : 0;
-            src1_p += swapped ? 0 : nc;
-        }
-
-        unary_gated_cuda<op>(src0_p, src1_p, (nv_bfloat16 *)dst_d, ggml_nelements(dst), nc, src0_o / sizeof(nv_bfloat16), src1_o / sizeof(nv_bfloat16), stream);
    } else {
        float * src0_p = (float *) src0_d;
        float * src1_p = (float *) src1_d;
@@ -674,7 +664,7 @@ static void ggml_cuda_op_unary_mul_impl(ggml_backend_cuda_context & ctx, ggml_te
    GGML_ASSERT(other_src->nb[0] == ggml_element_size(other_src));
    GGML_ASSERT(ggml_are_same_shape(unary_src, other_src));

-    GGML_ASSERT(unary_src->type == GGML_TYPE_F32 || unary_src->type == GGML_TYPE_F16 || unary_src->type == GGML_TYPE_BF16);
+    GGML_ASSERT(unary_src->type == GGML_TYPE_F32 || unary_src->type == GGML_TYPE_F16);
    GGML_ASSERT(unary_src->type == other_src->type);
    GGML_ASSERT(unary_src->type == mul_node->type);

@@ -689,10 +679,6 @@ static void ggml_cuda_op_unary_mul_impl(ggml_backend_cuda_context & ctx, ggml_te
        unary_gated_cuda<op>((const half *) unary_src->data, (const half *) other_src->data,
                             (half *) mul_node->data, k, nc,
                             unary_stride / sizeof(half), other_stride / sizeof(half), stream);
-    } else if (unary_src->type == GGML_TYPE_BF16) {
-        unary_gated_cuda<op>((const nv_bfloat16 *) unary_src->data, (const nv_bfloat16 *) other_src->data,
-                             (nv_bfloat16 *) mul_node->data, k, nc,
-                             unary_stride / sizeof(nv_bfloat16), other_stride / sizeof(nv_bfloat16), stream);
    } else {
        unary_gated_cuda<op>((const float *) unary_src->data, (const float *) other_src->data,
                             (float *) mul_node->data, k, nc,
@@ -471,7 +471,7 @@ static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmvq(

        const int vih = ((vh >> i) << 2) & 0x04040404;

-        const int vi = __vsub4(vil, vih);
+        const int vi = __vsubss4(vil, vih);

        sumf += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * sc); // SIMD dot product
    }
@@ -638,7 +638,7 @@ static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmvq(

        const int vih = ((vh >> (4*i)) << 4) & 0x30303030;

-        const int vi = __vsub4((vil | vih), 0x20202020); // vi = (vil | vih) - 32
+        const int vi = __vsubss4((vil | vih), 0x20202020); // vi = (vil | vih) - 32

        sumf += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * sc); // SIMD dot product
    }
@@ -5,11 +5,6 @@
 #include <mublas.h>
 #include <musa_bf16.h>
 #include <musa_fp16.h>
-
-#ifdef __MUSA_ARCH__
-#define __CUDA_ARCH__ 1300 // GGML_CUDA_CC_RUBIN
-#endif // __MUSA_ARCH__
-
 #define CUBLAS_COMPUTE_16F CUDA_R_16F
 #define CUBLAS_COMPUTE_32F CUDA_R_32F
 #define CUBLAS_COMPUTE_32F_FAST_16F MUBLAS_COMPUTE_32F_FAST_16F
@@ -15,7 +15,6 @@
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
-#include <string>
 #include <vector>

 #if __has_include(<filesystem>)
@@ -33,7 +32,9 @@ namespace fs = std::experimental::filesystem;
 * @brief prints the metadata of a single tensorf
 */
 static void ggml_et_dump_tensor_metadata(const ggml_tensor * ggtensor, size_t indent_level, const char * title) {
-    std::string spaces(indent_level, ' ');
+    char * spaces = (char *) alloca(indent_level + 1);
+    memset(spaces, ' ', indent_level);
+    spaces[indent_level] = '\0';
    fprintf(stderr,
            "%s%s: %s\n"
            "%s  type: %s\n"
@@ -42,10 +43,10 @@ static void ggml_et_dump_tensor_metadata(const ggml_tensor * ggtensor, size_t in
            "%s  op: %s\n"
            "%s  data: %p\n"
            "%s  src0: %p\n",
-            spaces.c_str(), title, ggtensor->name, spaces.c_str(), ggml_type_name(ggtensor->type), spaces.c_str(), (long long) ggtensor->ne[0],
-            (long long) ggtensor->ne[1], (long long) ggtensor->ne[2], (long long) ggtensor->ne[3], spaces.c_str(),
-            ggtensor->nb[0], ggtensor->nb[1], ggtensor->nb[2], ggtensor->nb[3], spaces.c_str(), ggml_op_name(ggtensor->op),
-            spaces.c_str(), ggtensor->data, spaces.c_str(), (void *) ggtensor->src[0]);
+            spaces, title, ggtensor->name, spaces, ggml_type_name(ggtensor->type), spaces, (long long) ggtensor->ne[0],
+            (long long) ggtensor->ne[1], (long long) ggtensor->ne[2], (long long) ggtensor->ne[3], spaces,
+            ggtensor->nb[0], ggtensor->nb[1], ggtensor->nb[2], ggtensor->nb[3], spaces, ggml_op_name(ggtensor->op),
+            spaces, ggtensor->data, spaces, (void *) ggtensor->src[0]);
 }

 /*
@@ -99,11 +99,10 @@ static int    opt_profile = 0; // profiling mode (0-disabled, 1-basic, 2-pmu)
 static bool   opt_hostbuf = false;
 static bool   opt_dma64   = false;

-static int    opt_mm_select  = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
-static int    opt_fa_select  = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
+static int    opt_mm_select = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
+static int    opt_fa_select = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
 static int    opt_gdn_select = 2; // 2 = HMX -> HVX, 1 = HVX, 0 = CPU (unsupported)
-static int    opt_ar_select  = 2; // 2 = fused ALLREDUCE+ADD (default), 1 = unfused ALLREDUCE, 0 = fallback to CPY+FENCE
-static int    opt_ar_scatter = 1; // 1 = reduce-scatter the fused ALLREDUCE+ADD (default), 0 = full reduction
+static int    opt_ar_select = 2; // 2 = fused ALLREDUCE+ADD (DMA, default), 1 = unfused ALLREDUCE (DMA), 0 = fallback to CPY+FENCE

 // Default PMU events, if profiling with PMU (mode=2) is enabled
 // See https://docs.qualcomm.com/doc/80-N2040-60/topic/pmu-events.html
@@ -401,7 +400,6 @@ static bool ggml_hexagon_precompute_allreduce_params(
    uint32_t n_ranks,
    bool has_add,
    bool is_row_bcast,
-    bool is_shard_ok,
    struct htp_allreduce_kernel_params * kparams
 );

@@ -2797,33 +2795,22 @@ struct ggml_hexagon_opbatch {
            return false;
        }

-        // scatter is only valid for in-place add within max outputs
-        const bool is_shard_ok = n_ranks <= HTP_OP_MAX_OUTPUTS && add_dst->data == ar_local->data;
+        for (uint32_t r = 0; r < n_ranks; r++) {
+            const ggml_tensor * ar_src = last_node.inputs[r];
+            if (ggml_hexagon_tensors_overlap(add_dst, ar_src)) {
+                HEX_VERBOSE("ggml-hex: %s skip ALLREDUCE_ADD fusion: dst overlaps allreduce src %u\n", sess->c_name(), r);
+                return false;
+            }
+        }

        struct htp_allreduce_kernel_params new_kparams;
        if (!ggml_hexagon_precompute_allreduce_params(
-            sess, add_dst, (uint32_t) ar_kparams->rank, (uint32_t) ar_kparams->n_ranks, true, is_row_bcast, is_shard_ok, &new_kparams
+            sess, add_dst, (uint32_t) ar_kparams->rank, (uint32_t) ar_kparams->n_ranks, true, is_row_bcast, &new_kparams
        )) {
            HEX_VERBOSE("ggml-hex: %s skip ALLREDUCE_ADD fusion: solver failed\n", sess->c_name());
            return false;
        }

-        const bool scatter_ok = new_kparams.n_dsts > 1;
-        new_kparams.mode = scatter_ok ? HTP_ALLREDUCE_SHARDED_FANOUT : HTP_ALLREDUCE_FULL;
-
-        for (uint32_t r = 0; r < n_ranks; r++) {
-            const ggml_tensor * ar_src = last_node.inputs[r];
-            if (!ggml_hexagon_tensors_overlap(add_dst, ar_src)) {
-                continue;
-            }
-            // in-place aliasing is safe under reduce-scatter since each rank writes disjoint shards
-            if (scatter_ok && r == rank) {
-                continue;
-            }
-            HEX_VERBOSE("ggml-hex: %s skip ALLREDUCE_ADD fusion: dst overlaps allreduce src %u\n", sess->c_name(), r);
-            return false;
-        }
-
        if (!try_fuse_common(res_tensor, add_dst)) {
            return false;
        }
@@ -2841,25 +2828,12 @@ struct ggml_hexagon_opbatch {
        memcpy(o.kernel_params, &new_kparams, sizeof(new_kparams));

        o.src[2 * n_ranks] = add_tensor(res_tensor);
-        if (new_kparams.mode == HTP_ALLREDUCE_SHARDED_FANOUT) {
-            // fan out shard to all per-rank partial buffers
-            GGML_ASSERT((uint32_t) new_kparams.n_dsts == n_ranks);
-            for (uint32_t d = 0; d < n_ranks; d++) {
-                GGML_ASSERT(o.src[d] != 0xffff);
-                o.dst[d] = o.src[d];
-            }
-            for (uint32_t d = n_ranks; d < HTP_OP_MAX_OUTPUTS; d++) {
-                o.dst[d] = 0xffff;
-            }
-        } else {
-            o.dst[0] = add_tensor(add_dst);
-            for (uint32_t d = 1; d < HTP_OP_MAX_OUTPUTS; d++) {
-                o.dst[d] = 0xffff;
-            }
+        o.dst[0]           = add_tensor(add_dst);
+        for (uint32_t d = 1; d < HTP_OP_MAX_OUTPUTS; d++) {
+            o.dst[d] = 0xffff;
        }

-        HEX_VERBOSE("ggml-hex: %s fused ALLREDUCE+ADD (#%u) mode=%d n_dsts=%d\n",
-                    sess->c_name(), n_ops - 1, (int) new_kparams.mode, (int) new_kparams.n_dsts);
+        HEX_VERBOSE("ggml-hex: %s fused ALLREDUCE+ADD (#%u)\n", sess->c_name(), n_ops - 1);
        return true;
    }

@@ -3800,7 +3774,6 @@ static bool ggml_hexagon_precompute_allreduce_params(
    uint32_t n_ranks,
    bool has_add,
    bool is_row_bcast,
-    bool is_shard_ok,
    struct htp_allreduce_kernel_params * kparams
 ) {
    memset(kparams, 0, sizeof(*kparams));
@@ -3820,20 +3793,11 @@ static bool ggml_hexagon_precompute_allreduce_params(
    const bool use_1d = is_contiguous && !(has_add && is_row_bcast && ne1 > 1);

    if (has_add) {
-        // sharded reduce-scatter for contiguous in-place add
-        if (opt_ar_scatter && is_shard_ok && use_1d && !is_row_bcast) {
-            const uint32_t rank_chunk_elems = hex_round_up((nelem + n_ranks - 1) / n_ranks, 128);
-            const uint32_t rank_elem_start  = (std::min)(rank * rank_chunk_elems, nelem);
-            const uint32_t rank_elem_end    = (std::min)(rank_elem_start + rank_chunk_elems, nelem);
-            kparams->n_dsts          = (int32_t) n_ranks;
-            kparams->rank_elem_start = (int32_t) rank_elem_start;
-            kparams->rank_nelem      = (int32_t) (rank_elem_end - rank_elem_start);
-        } else if (use_1d) {
-            kparams->n_dsts          = 1;
+        kparams->n_dsts = 1;
+        if (use_1d) {
            kparams->rank_elem_start = 0;
            kparams->rank_nelem      = (int32_t) nelem;
        } else {
-            kparams->n_dsts          = 1;
            kparams->rank_elem_start = 0;
            kparams->rank_nelem      = (int32_t) ne1;
        }
@@ -3856,8 +3820,6 @@ static bool ggml_hexagon_precompute_allreduce_params(
        }
    }

-    kparams->mode = (kparams->n_dsts > 1) ? HTP_ALLREDUCE_SHARDED_FANOUT : HTP_ALLREDUCE_FULL;
-
    if (use_1d) {
        const uint32_t rank_nelem = (uint32_t) kparams->rank_nelem;
        const uint32_t n_threads  = (std::min)((uint32_t) sess->n_threads, (std::max)(1u, rank_nelem / 128));
@@ -3963,8 +3925,9 @@ void ggml_hexagon_session::enqueue_allreduce(
    }

    ggml_hexagon_precompute_allreduce_params(
-        this, dst, rank, n_ranks, false, false, /*is_shard_ok=*/ false,
-        (struct htp_allreduce_kernel_params *) ar_node.kernel_params);
+        this, dst, rank, n_ranks, false, false,
+        (struct htp_allreduce_kernel_params *) ar_node.kernel_params
+    );

    ar_node.name = "ALLREDUCE";
    this->enqueue_op(ar_node);
@@ -6074,10 +6037,10 @@ static bool ggml_hexagon_supported_activations(const struct ggml_hexagon_session
    const struct ggml_tensor * src1 = op->src[1];
    const struct ggml_tensor * dst  = op;

-    if (src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_F16) {
+    if (src0->type != GGML_TYPE_F32) {
        return false;
    }
-    if (dst->type != src0->type) {
+    if (dst->type != GGML_TYPE_F32) {
        return false;
    }

@@ -6089,7 +6052,7 @@ static bool ggml_hexagon_supported_activations(const struct ggml_hexagon_session
    }

    if (src1) {
-        if (src1->type != src0->type) {
+        if (src1->type != GGML_TYPE_F32) {
            return false;
        }
        if (!ggml_are_same_shape(src0, src1)) {
@@ -6658,7 +6621,6 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) {
                case GGML_UNARY_OP_SILU:       return HTP_OP_UNARY_SILU;
                case GGML_UNARY_OP_GELU:       return HTP_OP_UNARY_GELU;
                case GGML_UNARY_OP_GELU_QUICK: return HTP_OP_UNARY_GELU;
-                case GGML_UNARY_OP_GELU_ERF:   return HTP_OP_UNARY_GELU_ERF;
                case GGML_UNARY_OP_SIGMOID:    return HTP_OP_UNARY_SIGMOID;
                case GGML_UNARY_OP_NEG:        return HTP_OP_UNARY_NEG;
                case GGML_UNARY_OP_EXP:        return HTP_OP_UNARY_EXP;
@@ -6674,12 +6636,11 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) {

        case GGML_OP_GLU:
            switch (ggml_get_glu_op(t)) {
-                case GGML_GLU_OP_SWIGLU:       return HTP_OP_GLU_SWIGLU;
-                case GGML_GLU_OP_SWIGLU_OAI:   return HTP_OP_GLU_SWIGLU_OAI;
+                case GGML_GLU_OP_SWIGLU:     return HTP_OP_GLU_SWIGLU;
+                case GGML_GLU_OP_SWIGLU_OAI: return HTP_OP_GLU_SWIGLU_OAI;
                case GGML_GLU_OP_SWIGLU_CLAMP: return HTP_OP_GLU_SWIGLU_CLAMP;
-                case GGML_GLU_OP_GEGLU:        return HTP_OP_GLU_GEGLU;
-                case GGML_GLU_OP_GEGLU_QUICK:  return HTP_OP_GLU_GEGLU_QUICK;
-                case GGML_GLU_OP_GEGLU_ERF:    return HTP_OP_GLU_GEGLU_ERF;
+                case GGML_GLU_OP_GEGLU:      return HTP_OP_GLU_GEGLU;
+                case GGML_GLU_OP_GEGLU_QUICK: return HTP_OP_GLU_GEGLU_QUICK;
                default: break;
            }
            break;
@@ -7736,7 +7697,6 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
                case GGML_UNARY_OP_SILU:
                case GGML_UNARY_OP_GELU:
                case GGML_UNARY_OP_GELU_QUICK:
-                case GGML_UNARY_OP_GELU_ERF:
                case GGML_UNARY_OP_RELU:
                case GGML_UNARY_OP_STEP:
                    supp = ggml_hexagon_supported_unary(sess, op);
@@ -7754,7 +7714,6 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
                case GGML_GLU_OP_SWIGLU_CLAMP:
                case GGML_GLU_OP_GEGLU:
                case GGML_GLU_OP_GEGLU_QUICK:
-                case GGML_GLU_OP_GEGLU_ERF:
                    supp = ggml_hexagon_supported_activations(sess, op);
                    break;
                default:
@@ -8036,7 +7995,7 @@ static bool ggml_backend_hexagon_comm_allreduce_tensor(void * comm_ctx_v, struct
    for (size_t r = 0; r < n_backends; r++) {
        auto sess = static_cast<ggml_hexagon_session *>(comm_ctx->backends[r]->context);
        struct htp_allreduce_kernel_params kparams;
-        if (!ggml_hexagon_precompute_allreduce_params(sess, tensors[r], (uint32_t) r, (uint32_t) n_backends, false, false, /*is_shard_ok=*/ false, &kparams)) {
+        if (!ggml_hexagon_precompute_allreduce_params(sess, tensors[r], (uint32_t) r, (uint32_t) n_backends, false, false, &kparams)) {
            return false;
        }
    }
@@ -8245,7 +8204,6 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
    const char * str_fa_select = getenv("GGML_HEXAGON_FA_SELECT");
    const char * str_gdn_select = getenv("GGML_HEXAGON_GDN_SELECT");
    const char * str_ar_select = getenv("GGML_HEXAGON_AR_SELECT");
-    const char * str_ar_scatter = getenv("GGML_HEXAGON_AR_SCATTER");
    const char * str_ndev     = getenv("GGML_HEXAGON_NDEV");
    const char * str_arch     = getenv("GGML_HEXAGON_ARCH");
    const char * str_vmem     = getenv("GGML_HEXAGON_VMEM");
@@ -8299,7 +8257,6 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
    opt_fa_select = str_fa_select ? atoi(str_fa_select)                   : opt_fa_select;
    opt_gdn_select = str_gdn_select ? atoi(str_gdn_select)                 : opt_gdn_select;
    opt_ar_select = str_ar_select ? atoi(str_ar_select)                   : opt_ar_select;
-    opt_ar_scatter = str_ar_scatter ? atoi(str_ar_scatter)                : opt_ar_scatter;
    opt_mbuf      = str_mbuf     ? strtoul(str_mbuf, NULL, 0) * MiB       : opt_mbuf;
    opt_vmem      = str_vmem     ? strtoul(str_vmem, NULL, 0) * MiB       : opt_vmem;
    opt_hostbuf   = str_hostbuf  ? atoi(str_hostbuf) != 0                 : opt_hostbuf;
@@ -55,22 +55,16 @@

 struct htp_act_context;

-typedef void (*glu_compute_f32_fn_t)(const float * restrict src0,
-                                     const float * restrict src1,
-                                     float * restrict dst,
-                                     const uint32_t num_rows,
-                                     const struct htp_act_context * actx);
-
-typedef void (*glu_compute_f16_fn_t)(const _Float16 * restrict src0,
-                                     const _Float16 * restrict src1,
-                                     _Float16 * restrict dst,
-                                     const uint32_t num_rows,
-                                     const struct htp_act_context * actx);
+typedef void (*glu_compute_fn_t)(const float * restrict src0,
+                                 const float * restrict src1,
+                                 float * restrict dst,
+                                 const uint32_t num_rows,
+                                 const struct htp_act_context * actx);

 struct htp_act_context {
    struct htp_ops_context * octx;

-    void *                   compute;
+    glu_compute_fn_t         compute;
    const char *             op_str;

    // Precomputed values
@@ -368,88 +362,6 @@ static inline void hvx_geglu_quick_f32_aa(uint8_t * restrict dst, const uint8_t
    }
 }

-static inline void hvx_geglu_erf_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
-    assert((unsigned long) dst  % 128 == 0);
-    assert((unsigned long) src0 % 128 == 0);
-    assert((unsigned long) src1 % 128 == 0);
-
-    HVX_Vector * restrict vdst        = (HVX_Vector *) dst;
-    const HVX_Vector * restrict vsrc0 = (const HVX_Vector *) src0;
-    const HVX_Vector * restrict vsrc1 = (const HVX_Vector *) src1;
-
-    const uint32_t epv  = 128 / sizeof(float);
-    const uint32_t nvec = n / epv;
-    const uint32_t nloe = n % epv;
-
-    uint32_t i = 0;
-
-    _Pragma("unroll(4)")
-    for (; i < nvec; i++) {
-        HVX_Vector x = vsrc0[i];
-        HVX_Vector g = vsrc1[i];
-        vdst[i] = hvx_vec_mul_f32_f32(hvx_vec_gelu_erf_f32(x), g);
-    }
-
-    if (nloe) {
-        HVX_Vector x = vsrc0[i];
-        HVX_Vector g = vsrc1[i];
-        HVX_Vector result = hvx_vec_mul_f32_f32(hvx_vec_gelu_erf_f32(x), g);
-        hvx_vec_store_a((void *) &vdst[i], nloe * sizeof(float), result);
-    }
-}
-
-static inline void hvx_geglu_f16_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
-    assert((unsigned long) dst  % 128 == 0);
-    assert((unsigned long) src0 % 128 == 0);
-    assert((unsigned long) src1 % 128 == 0);
-
-    HVX_Vector * restrict vdst        = (HVX_Vector *) dst;
-    const HVX_Vector * restrict vsrc0 = (const HVX_Vector *) src0;
-    const HVX_Vector * restrict vsrc1 = (const HVX_Vector *) src1;
-
-    const uint32_t epv  = 128 / sizeof(_Float16);
-    const uint32_t nvec = n / epv;
-    const uint32_t nloe = n % epv;
-
-    const HVX_Vector v_coef_a    = hvx_vec_splat_f16((_Float16) GELU_COEF_A);
-    const HVX_Vector v_sqrt_2_pi = hvx_vec_splat_f16((_Float16) SQRT_2_OVER_PI);
-    const HVX_Vector v_one       = hvx_vec_splat_f16((_Float16) 1.0f);
-    const HVX_Vector v_half      = hvx_vec_splat_f16((_Float16) 0.5f);
-
-    uint32_t i = 0;
-
-    _Pragma("unroll(4)")
-    for (; i < nvec; i++) {
-        HVX_Vector x = vsrc0[i];
-        HVX_Vector g = vsrc1[i];
-
-        // gelu(x) = 0.5*x*(1 + tanh(SQRT_2_OVER_PI*x*(1 + GELU_COEF_A*x*x)))
-        HVX_Vector x2    = hvx_vec_mul_f16_f16(x, x);
-        HVX_Vector coef  = hvx_vec_add_f16_f16(hvx_vec_mul_f16_f16(x2, v_coef_a), v_one);
-        HVX_Vector inner = hvx_vec_mul_f16_f16(hvx_vec_mul_f16_f16(x, coef), v_sqrt_2_pi);
-
-        HVX_Vector t = hvx_vec_add_f16_f16(hvx_vec_tanh_f16(inner), v_one);
-
-        HVX_Vector gelu_x = hvx_vec_mul_f16_f16(hvx_vec_mul_f16_f16(x, t), v_half);
-        vdst[i] = hvx_vec_mul_f16_f16(gelu_x, g);
-    }
-
-    if (nloe) {
-        HVX_Vector x = vsrc0[i];
-        HVX_Vector g = vsrc1[i];
-
-        HVX_Vector x2    = hvx_vec_mul_f16_f16(x, x);
-        HVX_Vector coef  = hvx_vec_add_f16_f16(hvx_vec_mul_f16_f16(x2, v_coef_a), v_one);
-        HVX_Vector inner = hvx_vec_mul_f16_f16(hvx_vec_mul_f16_f16(x, coef), v_sqrt_2_pi);
-
-        HVX_Vector t = hvx_vec_add_f16_f16(hvx_vec_tanh_f16(inner), v_one);
-
-        HVX_Vector gelu_x = hvx_vec_mul_f16_f16(hvx_vec_mul_f16_f16(x, t), v_half);
-        HVX_Vector res    = hvx_vec_mul_f16_f16(gelu_x, g);
-        hvx_vec_store_a((void *) &vdst[i], nloe * sizeof(_Float16), res);
-    }
-}
-
 // geglu(x, g) = gelu(x) * g
 static void geglu_f32(const float * restrict src0,
                                                const float * restrict src1,
@@ -484,207 +396,101 @@ static void geglu_quick_f32(const float * restrict src0,
    }
 }

-// geglu_erf(x, g) = gelu_erf(x) * g
-static void geglu_erf_f32(const float * restrict src0,
-                          const float * restrict src1,
-                          float * restrict dst,
-                          const uint32_t num_rows,
-                          const struct htp_act_context * actx) {
-    htp_glu_op_preamble;
+static void glu_f32_per_thread(unsigned int nth, unsigned int ith, void * data) {
+    struct htp_act_context * actx = (struct htp_act_context *) data;
+    htp_act_preamble;

-    for (uint32_t ib = 0; ib < num_rows; ib++) {
-        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
-        const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
-        uint8_t * restrict dst_ptr        = (uint8_t *) dst + (ib * dst_row_size_aligned);
+    struct htp_thread_trace * tr = actx->octx->ctx ? &actx->octx->ctx->trace[ith] : NULL;

-        hvx_geglu_erf_f32_aa(dst_ptr, src0_ptr, src1_ptr, nc);
+    size_t src0_row_size = actx->src0_row_size;
+    size_t src1_row_size = actx->src1_row_size;
+    size_t dst_row_size  = actx->dst_row_size;
+
+    size_t src0_row_stride = actx->src0_row_stride;
+    size_t src1_row_stride = actx->src1_row_stride;
+
+    const uint32_t src0_nrows            = actx->src0_nrows;
+    const uint32_t src0_nrows_per_thread = actx->src0_nrows_per_thread;
+
+    const uint32_t src0_start_row = actx->row_start + src0_nrows_per_thread * ith;
+    const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, actx->row_start + src0_nrows);
+
+    /* no work for this thread */
+    if (src0_start_row >= src0_end_row) {
+        return;
    }
+
+    const dma_addr_t data_src0 = actx->data_src0;
+    const dma_addr_t data_src1 = actx->data_src1;
+    const dma_addr_t data_dst  = actx->data_dst;
+
+    const size_t src0_row_size_aligned = actx->src0_row_size_aligned;
+    const size_t src1_row_size_aligned = actx->src1_row_size_aligned;
+    const size_t dst_row_size_aligned  = actx->dst_row_size_aligned;
+
+    uint8_t * restrict src0_spad_data = actx->vtcm_src0 + (ith * actx->vtcm_src0_size_per_thread);
+    uint8_t * restrict src1_spad_data = actx->vtcm_src1 + (ith * actx->vtcm_src1_size_per_thread);
+    uint8_t * restrict dst_spad_data  = actx->vtcm_dst  + (ith * actx->vtcm_dst_size_per_thread);
+
+    size_t src0_spad_half_size = actx->src0_spad_half_size;
+    size_t src1_spad_half_size = actx->src1_spad_half_size;
+    size_t dst_spad_half_size  = actx->dst_spad_half_size;
+
+    const int BLOCK = actx->block;
+    if (BLOCK == 0) {
+        FARF(ERROR, "%s : VTCM reservation %zu is too small, needed %zu\n",
+             actx->op_str, actx->vtcm_src0_size_per_thread, src0_row_size_aligned);
+        return;
+    }
+
+    dma_queue * dma_q = actx->octx->ctx->dma[ith];
+    glu_compute_fn_t compute = actx->compute;
+
+    for (uint32_t ir = src0_start_row, spad_idx = 0; ir < src0_end_row && spad_idx < 2; ir += BLOCK, spad_idx++) {
+        const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);
+
+        /* Dummy DMA transation for sequencing (interleaving dst,src,dst,...) */
+        dma_queue_push(dma_q,
+                       dma_make_data(data_dst, dst_spad_data + (spad_idx * dst_spad_half_size)),
+                       dst_row_size, dst_row_size_aligned, dst_row_size, 0);
+
+        dma_queue_push(dma_q,
+            dma_make_data(src0_spad_data + (spad_idx * src0_spad_half_size), data_src0 + (ir * src0_row_stride)),
+            src0_row_size_aligned, src0_row_stride, src0_row_size, block_size);
+
+        dma_queue_push(dma_q,
+            dma_make_data(src1_spad_data + (spad_idx * src1_spad_half_size), data_src1 + (ir * src1_row_stride)),
+            src1_row_size_aligned, src1_row_stride, src1_row_size, block_size);
+    }
+
+    for (uint32_t ir = src0_start_row; ir < src0_end_row; ir += BLOCK) {
+        const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);
+
+        float * dst_spad  = (float *) dma_queue_pop(dma_q).src;
+        float * src0_spad = (float *) dma_queue_pop(dma_q).dst;
+        float * src1_spad = (float *) dma_queue_pop(dma_q).dst;
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir);
+        compute(src0_spad, src1_spad, dst_spad, block_size, actx);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir);
+
+        dma_queue_push(dma_q, dma_make_data(data_dst + (ir * dst_row_size), dst_spad),
+                       dst_row_size, dst_row_size_aligned, dst_row_size, block_size);
+
+        /* prefetch N+2 loop iteration if any */
+        const uint32_t pref_block = (ir + BLOCK * 2);
+        if (pref_block < src0_end_row) {
+            const uint32_t pref_block_size = MIN(BLOCK, src0_end_row - pref_block);
+            dma_queue_push(dma_q, dma_make_data(src0_spad, data_src0 + (pref_block * src0_row_stride)),
+                           src0_row_size_aligned, src0_row_stride, src0_row_size, pref_block_size);
+            dma_queue_push(dma_q, dma_make_data(src1_spad, data_src1 + (pref_block * src1_row_stride)),
+                           src1_row_size_aligned, src1_row_stride, src1_row_size, pref_block_size);
+        }
+    }
+
+    dma_queue_flush(dma_q);
 }

-// swiglu(x) = x1 * sigmoid(x0)
-static void swiglu_f16(const _Float16 * restrict src0,
-                       const _Float16 * restrict src1,
-                       _Float16 * restrict dst,
-                       const uint32_t num_rows,
-                       const struct htp_act_context * actx) {
-    htp_glu_op_preamble;
-
-    for (uint32_t ib = 0; ib < num_rows; ib++) {
-        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
-        const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
-        uint8_t * restrict dst_ptr        = (uint8_t *) dst + (ib * dst_row_size_aligned);
-
-        hvx_sigmoid_f16_aa(dst_ptr, src0_ptr, nc);
-        hvx_mul_mul_f16_aa(dst_ptr, src0_ptr, dst_ptr, src1_ptr, nc);
-    }
-}
-
-// out = x * sigmoid(alpha * x) * (clamp(y, -limit, limit) + 1.f)
-static void swiglu_oai_f16(const _Float16 * restrict src0,
-                           const _Float16 * restrict src1,
-                           _Float16 * restrict dst,
-                           const uint32_t num_rows,
-                           const struct htp_act_context * actx) {
-    htp_glu_op_preamble;
-    const _Float16 alpha = (_Float16) ((const float *) (actx->octx->op_params))[2];
-    const _Float16 limit = (_Float16) ((const float *) (actx->octx->op_params))[3];
-
-    for (uint32_t ib = 0; ib < num_rows; ib++) {
-        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
-        const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
-        uint8_t * restrict dst_ptr        = (uint8_t *) dst + (ib * dst_row_size_aligned);
-
-        // x (src0_ptr) = std::min(src0_p[k], limit);
-        hvx_min_scalar_f16((uint8_t *) src0_ptr, src0_ptr, limit, nc);
-        // y1 (src1_ptr) = std::clamp(src1_p[k], -limit, limit);
-        hvx_clamp_scalar_f16((uint8_t *) src1_ptr, src1_ptr, -limit, limit, nc);
-        // y (src1_ptr) = y1 + 1.f
-        hvx_add_scalar_f16((uint8_t *) src1_ptr, src1_ptr, (_Float16) 1.0f, nc);
-        // x1 (dst_ptr) = alpha * x
-        hvx_mul_scalar_f16(dst_ptr, src0_ptr, alpha, nc);
-        // x2 (dst_ptr) = sigmoid(x1) = 1/(1+exp(-x1))
-        hvx_sigmoid_f16_aa(dst_ptr, dst_ptr, nc);
-        // out = x * sigmoid(alpha * x) * (y + 1.f)
-        hvx_mul_mul_f16_aa(dst_ptr, src0_ptr, dst_ptr, src1_ptr, nc);
-    }
-}
-
-static void swiglu_clamp_f16(const _Float16 * restrict src0,
-                             const _Float16 * restrict src1,
-                             _Float16 * restrict dst,
-                             const uint32_t                 num_rows,
-                             const struct htp_act_context * actx) {
-    htp_glu_op_preamble;
-    const _Float16 limit = (_Float16) ((const float *) (actx->octx->op_params))[3];
-
-    for (uint32_t ib = 0; ib < num_rows; ib++) {
-        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
-        const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
-        uint8_t * restrict dst_ptr        = (uint8_t *) dst + (ib * dst_row_size_aligned);
-
-        hvx_min_scalar_f16((uint8_t *) src0_ptr, src0_ptr, limit, nc);
-        hvx_clamp_scalar_f16((uint8_t *) src1_ptr, src1_ptr, -limit, limit, nc);
-        hvx_sigmoid_f16_aa(dst_ptr, src0_ptr, nc);
-        hvx_mul_mul_f16_aa(dst_ptr, src0_ptr, dst_ptr, src1_ptr, nc);
-    }
-}
-
-// geglu(x, g) = gelu(x) * g, gelu via tanh approximation
-static void geglu_f16(const _Float16 * restrict src0,
-                      const _Float16 * restrict src1,
-                      _Float16 * restrict dst,
-                      const uint32_t num_rows,
-                      const struct htp_act_context * actx) {
-    htp_glu_op_preamble;
-
-    for (uint32_t ib = 0; ib < num_rows; ib++) {
-        const uint8_t * restrict src0_ptr = (const uint8_t *) src0 + (ib * src0_row_size_aligned);
-        const uint8_t * restrict src1_ptr = (const uint8_t *) src1 + (ib * src1_row_size_aligned);
-        uint8_t * restrict dst_ptr        = (uint8_t *) dst + (ib * dst_row_size_aligned);
-
-        hvx_geglu_f16_aa(dst_ptr, src0_ptr, src1_ptr, nc);
-    }
-}
-
-#define DEFINE_GLU_PER_THREAD(TYPE, SUFFIX)                                                                            \
-    static void glu_##SUFFIX##_per_thread(unsigned int nth, unsigned int ith, void * data) {                           \
-        struct htp_act_context * actx = (struct htp_act_context *) data;                                               \
-        htp_act_preamble;                                                                                              \
-                                                                                                                       \
-        struct htp_thread_trace * tr = actx->octx->ctx ? &actx->octx->ctx->trace[ith] : NULL;                          \
-                                                                                                                       \
-        size_t src0_row_size = actx->src0_row_size;                                                                    \
-        size_t src1_row_size = actx->src1_row_size;                                                                    \
-        size_t dst_row_size  = actx->dst_row_size;                                                                     \
-                                                                                                                       \
-        size_t src0_row_stride = actx->src0_row_stride;                                                                \
-        size_t src1_row_stride = actx->src1_row_stride;                                                                \
-                                                                                                                       \
-        const uint32_t src0_nrows            = actx->src0_nrows;                                                       \
-        const uint32_t src0_nrows_per_thread = actx->src0_nrows_per_thread;                                            \
-                                                                                                                       \
-        const uint32_t src0_start_row = actx->row_start + src0_nrows_per_thread * ith;                                 \
-        const uint32_t src0_end_row   = MIN(src0_start_row + src0_nrows_per_thread, actx->row_start + src0_nrows);     \
-                                                                                                                       \
-        /* no work for this thread */                                                                                  \
-        if (src0_start_row >= src0_end_row) {                                                                          \
-            return;                                                                                                    \
-        }                                                                                                              \
-                                                                                                                       \
-        const dma_addr_t data_src0 = actx->data_src0;                                                                  \
-        const dma_addr_t data_src1 = actx->data_src1;                                                                  \
-        const dma_addr_t data_dst  = actx->data_dst;                                                                   \
-                                                                                                                       \
-        const size_t src0_row_size_aligned = actx->src0_row_size_aligned;                                              \
-        const size_t src1_row_size_aligned = actx->src1_row_size_aligned;                                              \
-        const size_t dst_row_size_aligned  = actx->dst_row_size_aligned;                                               \
-                                                                                                                       \
-        uint8_t * restrict src0_spad_data = actx->vtcm_src0 + (ith * actx->vtcm_src0_size_per_thread);                 \
-        uint8_t * restrict src1_spad_data = actx->vtcm_src1 + (ith * actx->vtcm_src1_size_per_thread);                 \
-        uint8_t * restrict dst_spad_data  = actx->vtcm_dst  + (ith * actx->vtcm_dst_size_per_thread);                  \
-                                                                                                                       \
-        size_t src0_spad_half_size = actx->src0_spad_half_size;                                                        \
-        size_t src1_spad_half_size = actx->src1_spad_half_size;                                                        \
-        size_t dst_spad_half_size  = actx->dst_spad_half_size;                                                         \
-                                                                                                                       \
-        const int BLOCK = actx->block;                                                                                 \
-        if (BLOCK == 0) {                                                                                              \
-            FARF(ERROR, "%s : VTCM reservation %zu is too small, needed %zu\n",                                        \
-                 actx->op_str, actx->vtcm_src0_size_per_thread, src0_row_size_aligned);                                \
-            return;                                                                                                    \
-        }                                                                                                              \
-                                                                                                                       \
-        dma_queue * dma_q = actx->octx->ctx->dma[ith];                                                                 \
-        glu_compute_##SUFFIX##_fn_t compute = (glu_compute_##SUFFIX##_fn_t) actx->compute;                             \
-                                                                                                                       \
-        for (uint32_t ir = src0_start_row, spad_idx = 0; ir < src0_end_row && spad_idx < 2; ir += BLOCK, spad_idx++) { \
-            const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);                                                 \
-                                                                                                                       \
-            /* Dummy DMA transation for sequencing (interleaving dst,src,dst,...) */                                   \
-            dma_queue_push(dma_q,                                                                                      \
-                           dma_make_data(data_dst, dst_spad_data + (spad_idx * dst_spad_half_size)),                   \
-                           dst_row_size, dst_row_size_aligned, dst_row_size, 0);                                       \
-                                                                                                                       \
-            dma_queue_push(dma_q,                                                                                      \
-                dma_make_data(src0_spad_data + (spad_idx * src0_spad_half_size), data_src0 + (ir * src0_row_stride)),  \
-                src0_row_size_aligned, src0_row_stride, src0_row_size, block_size);                                    \
-                                                                                                                       \
-            dma_queue_push(dma_q,                                                                                      \
-                dma_make_data(src1_spad_data + (spad_idx * src1_spad_half_size), data_src1 + (ir * src1_row_stride)),  \
-                src1_row_size_aligned, src1_row_stride, src1_row_size, block_size);                                    \
-        }                                                                                                              \
-                                                                                                                       \
-        for (uint32_t ir = src0_start_row; ir < src0_end_row; ir += BLOCK) {                                           \
-            const uint32_t block_size = MIN(BLOCK, src0_end_row - ir);                                                 \
-                                                                                                                       \
-            TYPE * dst_spad  = (TYPE *) dma_queue_pop(dma_q).src;                                                      \
-            TYPE * src0_spad = (TYPE *) dma_queue_pop(dma_q).dst;                                                      \
-            TYPE * src1_spad = (TYPE *) dma_queue_pop(dma_q).dst;                                                      \
-                                                                                                                       \
-            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir);                                                     \
-            compute(src0_spad, src1_spad, dst_spad, block_size, actx);                                                 \
-            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir);                                                      \
-                                                                                                                       \
-            dma_queue_push(dma_q, dma_make_data(data_dst + (ir * dst_row_size), dst_spad),                             \
-                           dst_row_size, dst_row_size_aligned, dst_row_size, block_size);                              \
-                                                                                                                       \
-            /* prefetch N+2 loop iteration if any */                                                                   \
-            const uint32_t pref_block = (ir + BLOCK * 2);                                                              \
-            if (pref_block < src0_end_row) {                                                                           \
-                const uint32_t pref_block_size = MIN(BLOCK, src0_end_row - pref_block);                                \
-                dma_queue_push(dma_q, dma_make_data(src0_spad, data_src0 + (pref_block * src0_row_stride)),            \
-                               src0_row_size_aligned, src0_row_stride, src0_row_size, pref_block_size);                \
-                dma_queue_push(dma_q, dma_make_data(src1_spad, data_src1 + (pref_block * src1_row_stride)),            \
-                               src1_row_size_aligned, src1_row_stride, src1_row_size, pref_block_size);                \
-            }                                                                                                          \
-        }                                                                                                              \
-                                                                                                                       \
-        dma_queue_flush(dma_q);                                                                                        \
-    }
-
-DEFINE_GLU_PER_THREAD(float, f32)
-DEFINE_GLU_PER_THREAD(_Float16, f16)
-
 static int execute_op_activations_f32(struct htp_ops_context * octx) {
    const struct htp_tensor * src0 = octx->src[0];
    const struct htp_tensor * src1 = octx->src[1];
@@ -695,8 +501,8 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
        return HTP_STATUS_NO_SUPPORT;
    }

-    glu_compute_f32_fn_t compute_fn = NULL;
-    const char *         op_type    = NULL;
+    glu_compute_fn_t compute_fn = NULL;
+    const char *      op_type    = NULL;

    switch (octx->op) {
        case HTP_OP_GLU_SWIGLU:
@@ -723,12 +529,6 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
            compute_fn = geglu_quick_f32;
            op_type    = "geglu-quick-f32";
            break;
-
-        case HTP_OP_GLU_GEGLU_ERF:
-            compute_fn = geglu_erf_f32;
-            op_type    = "geglu-erf-f32";
-            break;
-
        default:
            FARF(ERROR, "Unsupported activations Op %u\n", octx->op);
            return HTP_STATUS_NO_SUPPORT;
@@ -791,7 +591,7 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
    // Prepare context
    struct htp_act_context actx;
    actx.octx    = octx;
-    actx.compute = (void *) compute_fn;
+    actx.compute = compute_fn;
    actx.op_str  = op_type;

    actx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
@@ -834,8 +634,7 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
                  octx->op == HTP_OP_GLU_SWIGLU_OAI ||
                  octx->op == HTP_OP_GLU_SWIGLU_CLAMP ||
                  octx->op == HTP_OP_GLU_GEGLU ||
-                  octx->op == HTP_OP_GLU_GEGLU_QUICK ||
-                  octx->op == HTP_OP_GLU_GEGLU_ERF)) {
+                  octx->op == HTP_OP_GLU_GEGLU_QUICK)) {
         const int32_t swapped = octx->op_params[1];
         data_src1 = data_src0;
         actx.src1_row_size = actx.src0_row_size;
@@ -856,174 +655,11 @@ static int execute_op_activations_f32(struct htp_ops_context * octx) {
    return HTP_STATUS_OK;
 }

-static int execute_op_activations_f16(struct htp_ops_context * octx) {
-    const struct htp_tensor * src0 = octx->src[0];
-    const struct htp_tensor * src1 = octx->src[1];
-    const struct htp_tensor * dst  = octx->dst;
-
-    if ((dst->ne[0] * SIZEOF_FP16) != dst->nb[1]) {
-        FARF(ERROR, "Non-contiguous dst is not supported at this time \n");
-        return HTP_STATUS_NO_SUPPORT;
-    }
-
-    glu_compute_f16_fn_t compute_fn = NULL;
-    const char *         op_type    = NULL;
-
-    switch (octx->op) {
-        case HTP_OP_GLU_SWIGLU:
-            compute_fn = swiglu_f16;
-            op_type    = "swiglu-f16";
-            break;
-
-        case HTP_OP_GLU_SWIGLU_OAI:
-            compute_fn = swiglu_oai_f16;
-            op_type    = "swiglu-oai-f16";
-            break;
-
-        case HTP_OP_GLU_SWIGLU_CLAMP:
-            compute_fn = swiglu_clamp_f16;
-            op_type    = "swiglu-clamp-f16";
-            break;
-
-        case HTP_OP_GLU_GEGLU:
-            compute_fn = geglu_f16;
-            op_type    = "geglu-f16";
-            break;
-
-        default:
-            FARF(ERROR, "Unsupported activations Op %u\n", octx->op);
-            return HTP_STATUS_NO_SUPPORT;
-    }
-
-    const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3];
-    const size_t dst_row_size = dst->ne[0] * SIZEOF_FP16;
-
-    uint32_t row_start = 0;
-    uint32_t nrows     = src0_nrows;
-
-    if (octx->ctx->mdev.count > 1) {
-        uint32_t rows_per_chunk = 0;
-        htp_tensor_mdev_rows_per_chunk(dst, sizeof(_Float16), (uint32_t) dst_row_size, &rows_per_chunk);
-        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(src0_nrows, rows_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 = octx->n_threads;
-
-    // row_size   = bytes of useful data per row (what the kernel touches / what DMA copies).
-    // row_stride = bytes between successive rows in DDR (may exceed row_size for non-contig src).
-    const size_t nc_bytes        = dst_row_size;
-    const size_t src0_row_size   = nc_bytes;
-    const size_t src1_row_size   = nc_bytes;
-    const size_t src0_row_stride = src0->nb[1];
-    const size_t src1_row_stride = src1 ? src1->nb[1] : src0->nb[1];
-
-    const size_t src0_row_size_aligned = hex_round_up(src0_row_size, VLEN);
-    const size_t src1_row_size_aligned = hex_round_up(src1_row_size, VLEN);
-    const size_t dst_row_size_aligned  = hex_round_up(dst_row_size, VLEN);
-
-    struct htp_act_vtcm_layout L;
-    htp_act_vtcm_layout_build(&L, src0_row_size_aligned, src1_row_size_aligned, dst_row_size_aligned, n_threads,
-                              octx->ctx->vtcm_size);
-
-    // Make sure the reserved vtcm size is sufficient
-    if (L.vtcm_row_per_thread == 0) {
-        FARF(ERROR, "act-%s : current VTCM reservation %zu is too small for even 1 row per thread, needed at least %zu\n", op_type, octx->ctx->vtcm_size,
-             (src0_row_size_aligned + src1_row_size_aligned + dst_row_size_aligned) * n_threads);
-        return HTP_STATUS_VTCM_TOO_SMALL;
-    }
-
-    if (src1) {
-        FARF(HIGH, "%s: %ux%ux%ux%u x %ux%ux%ux%u -> %ux%ux%ux%u : src0-vtcm-size %zu src1-vtcm-size %zu dst-vtcm-size %zu\n",
-             op_type, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], src1->ne[0], src1->ne[1], src1->ne[2],
-             src1->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], L.src0_bytes_per_thread * n_threads,
-             L.src1_bytes_per_thread * n_threads, L.dst_bytes_per_thread * n_threads);
-    } else {
-        FARF(HIGH, "%s: %ux%ux%ux%u -> %ux%ux%ux%u : src0-vtcm-size %zu src1-vtcm-size %zu dst-vtcm-size %zu\n", op_type,
-             src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
-             L.src0_bytes_per_thread * n_threads, L.src1_bytes_per_thread * n_threads, L.dst_bytes_per_thread * n_threads);
-    }
-
-    // Prepare context
-    struct htp_act_context actx;
-    actx.octx    = octx;
-    actx.compute = (void *) compute_fn;
-    actx.op_str  = op_type;
-
-    actx.src0_nrows_per_thread = fastdiv(nrows + n_threads - 1, &octx->n_threads_div);
-
-    actx.src0_row_size = src0_row_size;
-    actx.src1_row_size = src1_row_size;
-    actx.dst_row_size  = dst_row_size;
-
-    actx.src0_row_size_aligned = src0_row_size_aligned;
-    actx.src1_row_size_aligned = src1_row_size_aligned;
-    actx.dst_row_size_aligned  = dst_row_size_aligned;
-
-    actx.src0_row_stride = src0_row_stride;
-    actx.src1_row_stride = src1_row_stride;
-
-    uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base;
-    actx.vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0);
-    actx.vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1);
-    actx.vtcm_dst  = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst);
-
-    actx.vtcm_src0_size_per_thread = L.src0_bytes_per_thread;
-    actx.vtcm_src1_size_per_thread = L.src1_bytes_per_thread;
-    actx.vtcm_dst_size_per_thread  = L.dst_bytes_per_thread;
-
-    actx.src0_spad_half_size = L.src0_bytes_per_thread / 2;
-    actx.src1_spad_half_size = L.src1_bytes_per_thread / 2;
-    actx.dst_spad_half_size  = L.dst_bytes_per_thread / 2;
-
-    actx.block = actx.src0_spad_half_size / actx.src0_row_size_aligned;
-    actx.src0_nrows = nrows;
-    actx.row_start  = row_start;
-
-    actx.nc = dst->ne[0];
-
-    // Addresses and GLU logic
-    dma_addr_t data_src0 = src0->data;
-    dma_addr_t data_src1 = src1 ? src1->data : 0;
-
-    if (!src1 && (octx->op == HTP_OP_GLU_SWIGLU ||
-                  octx->op == HTP_OP_GLU_SWIGLU_OAI ||
-                  octx->op == HTP_OP_GLU_SWIGLU_CLAMP ||
-                  octx->op == HTP_OP_GLU_GEGLU)) {
-         const int32_t swapped = octx->op_params[1];
-         data_src1 = data_src0;
-         actx.src1_row_size = actx.src0_row_size;
-
-         size_t nc_in_bytes = actx.nc * SIZEOF_FP16;
-         if (swapped) {
-             data_src0 += nc_in_bytes;
-         } else {
-             data_src1 += nc_in_bytes;
-         }
-    }
-
-    actx.data_src0 = data_src0;
-    actx.data_src1 = data_src1;
-    actx.data_dst  = dst->data;
-
-    work_queue_run(octx->ctx->work_queue, (worker_callback_t)glu_f16_per_thread, &actx, n_threads);
-    return HTP_STATUS_OK;
-}
-
-
 int op_activations(struct htp_ops_context * octx) {
    switch (octx->src[0]->type) {
        case HTP_TYPE_F32:
            return execute_op_activations_f32(octx);

-        case HTP_TYPE_F16:
-            return execute_op_activations_f16(octx);
-
        default:
            return HTP_STATUS_NO_SUPPORT;
    }
@@ -38,103 +38,97 @@ struct htp_allreduce_context {
    uint8_t * res_spad_base;
 };

-#define DEFINE_ALLREDUCE_THREAD_DMA_1D(SUFFIX, TYPE, HVX_ADD_FN, HAS_ADD)                                         \
-static void allreduce_thread_dma_1d_##SUFFIX(unsigned int nth, unsigned int ith, void * data) {                   \
-    struct htp_allreduce_context * actx = (struct htp_allreduce_context *) data;                                  \
-    struct htp_ops_context * octx = actx->octx;                                                                   \
-                                                                                                                  \
-    const uint32_t n_ranks     = actx->n_ranks;                                                                   \
-    const uint32_t n_dsts      = actx->n_dsts;                                                                    \
-    const uint32_t block_elems = actx->block_elems;                                                               \
-                                                                                                                  \
-    const uint32_t dr  = actx->elems_per_thread;                                                                  \
-    const uint32_t ir0 = actx->rank_elem_start + dr * ith;                                                        \
-    const uint32_t ir1 = MIN(ir0 + dr, actx->rank_elem_start + actx->rank_nelem);                                 \
-    if (ir0 >= ir1) return;                                                                                       \
-                                                                                                                  \
-    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                        \
-    dma_queue * dma_q = octx->ctx->dma[ith];                                                                      \
-                                                                                                                  \
-    const size_t vtcm_thread_offset = ith * actx->vtcm_size_per_thread;                                           \
-    uint8_t * dst_spad_base = actx->dst_spad_base + vtcm_thread_offset;                                           \
-    uint8_t * res_spad_base = HAS_ADD ? (actx->res_spad_base + vtcm_thread_offset) : NULL;                        \
-                                                                                                                  \
-    const size_t spad_half = actx->vtcm_size_per_thread / 2;                                                      \
-    uint32_t ir_prefetch = ir0;                                                                                   \
-    int spad_idx = 0;                                                                                             \
-                                                                                                                  \
-    for (int k = 0; k < 2 && ir_prefetch < ir1; k++) {                                                            \
-        uint32_t cur_elems = MIN(block_elems, ir1 - ir_prefetch);                                                 \
-        size_t   cur_bytes = cur_elems * sizeof(TYPE);                                                            \
-        uint8_t * d_spad = dst_spad_base + spad_idx * spad_half;                                                  \
-        for (uint32_t d = 0; d < n_dsts; d++) {                                                                   \
-            dma_addr_t d_ddr = octx->dsts[d]->data + ir_prefetch * sizeof(TYPE);                                  \
-            dma_queue_push(dma_q, dma_make_data(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 0);              \
-        }                                                                                                         \
-        for (uint32_t s = 0; s < n_ranks; s++) {                                                                  \
-            uint8_t * s_spad = actx->src_spad_base[s] + vtcm_thread_offset + spad_idx * spad_half;                \
-            const dma_addr_t s_ddr = octx->src[s]->data + ir_prefetch * sizeof(TYPE);                             \
-            dma_queue_push(dma_q, dma_make_data(s_spad, s_ddr), cur_bytes, cur_bytes, cur_bytes, 1);              \
-        }                                                                                                         \
-        if (HAS_ADD) {                                                                                            \
-            uint8_t * r_spad = res_spad_base + spad_idx * spad_half;                                              \
-            const dma_addr_t r_ddr = octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE);                   \
-            dma_queue_push(dma_q, dma_make_data(r_spad, r_ddr), cur_bytes, cur_bytes, cur_bytes, 1);              \
-        }                                                                                                         \
-        ir_prefetch += cur_elems;                                                                                 \
-        spad_idx ^= 1;                                                                                            \
-    }                                                                                                             \
-                                                                                                                  \
-    int comp_spad_idx = 0;                                                                                        \
-    for (uint32_t ir = ir0; ir < ir1; ) {                                                                         \
-        uint32_t cur_elems = MIN(block_elems, ir1 - ir);                                                          \
-        size_t   cur_bytes = cur_elems * sizeof(TYPE);                                                            \
-        for (uint32_t d = 0; d < n_dsts; d++) {                                                                   \
-            dma_queue_pop(dma_q);                                                                                 \
-        }                                                                                                         \
-        for (uint32_t s = 0; s < n_ranks; s++) {                                                                  \
-            dma_queue_pop(dma_q);                                                                                 \
-        }                                                                                                         \
-        if (HAS_ADD) {                                                                                            \
-            dma_queue_pop(dma_q);                                                                                 \
-        }                                                                                                         \
-        uint8_t * d_spad  = dst_spad_base + comp_spad_idx * spad_half;                                            \
-        uint8_t * s0_spad = actx->src_spad_base[0] + vtcm_thread_offset + comp_spad_idx * spad_half;              \
-        uint8_t * s1_spad = actx->src_spad_base[1] + vtcm_thread_offset + comp_spad_idx * spad_half;              \
-        uint8_t * r_spad  = HAS_ADD ? (res_spad_base + comp_spad_idx * spad_half) : NULL;                         \
-        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);                                         \
-        HVX_ADD_FN(d_spad, s0_spad, s1_spad, cur_elems);                                                          \
-        for (uint32_t s = 2; s < n_ranks; s++) {                                                                  \
-            uint8_t * ss_spad = actx->src_spad_base[s] + vtcm_thread_offset + comp_spad_idx * spad_half;          \
-            HVX_ADD_FN(d_spad, d_spad, ss_spad, cur_elems);                                                       \
-        }                                                                                                         \
-        if (HAS_ADD) {                                                                                            \
-            HVX_ADD_FN(d_spad, d_spad, r_spad, cur_elems);                                                        \
-        }                                                                                                         \
-        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);                                          \
-        for (uint32_t d = 0; d < n_dsts; d++) {                                                                   \
-            dma_addr_t d_ddr = octx->dsts[d]->data + ir * sizeof(TYPE);                                           \
-            dma_queue_push(dma_q, dma_make_data(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 1);              \
-        }                                                                                                         \
-        if (ir_prefetch < ir1) {                                                                                  \
-            uint32_t next_elems = MIN(block_elems, ir1 - ir_prefetch);                                            \
-            size_t   next_bytes = next_elems * sizeof(TYPE);                                                      \
-            for (uint32_t s = 0; s < n_ranks; s++) {                                                              \
-                uint8_t * s_spad = actx->src_spad_base[s] + vtcm_thread_offset + comp_spad_idx * spad_half;       \
-                const dma_addr_t s_next = octx->src[s]->data + ir_prefetch * sizeof(TYPE);                        \
-                dma_queue_push(dma_q, dma_make_data(s_spad, s_next), next_bytes, next_bytes, next_bytes, 1);      \
-            }                                                                                                     \
-            if (HAS_ADD) {                                                                                        \
-                uint8_t * r_spad_next = res_spad_base + comp_spad_idx * spad_half;                                \
-                const dma_addr_t r_next = octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE);              \
-                dma_queue_push(dma_q, dma_make_data(r_spad_next, r_next), next_bytes, next_bytes, next_bytes, 1); \
-            }                                                                                                     \
-            ir_prefetch += next_elems;                                                                            \
-        }                                                                                                         \
-        comp_spad_idx ^= 1;                                                                                       \
-        ir += cur_elems;                                                                                          \
-    }                                                                                                             \
-    dma_queue_flush(dma_q);                                                                                       \
+#define DEFINE_ALLREDUCE_THREAD_DMA_1D(SUFFIX, TYPE, HVX_ADD_FN, HAS_ADD)                                       \
+static void allreduce_thread_dma_1d_##SUFFIX(unsigned int nth, unsigned int ith, void * data) {                 \
+    struct htp_allreduce_context * actx = (struct htp_allreduce_context *) data;                                \
+    struct htp_ops_context * octx = actx->octx;                                                                 \
+                                                                                                                \
+    const uint32_t n_ranks     = actx->n_ranks;                                                                 \
+    const uint32_t n_dsts      = actx->n_dsts;                                                                  \
+    const uint32_t block_elems = actx->block_elems;                                                             \
+                                                                                                                \
+    const uint32_t dr  = actx->elems_per_thread;                                                                \
+    const uint32_t ir0 = actx->rank_elem_start + dr * ith;                                                      \
+    const uint32_t ir1 = MIN(ir0 + dr, actx->rank_elem_start + actx->rank_nelem);                               \
+    if (ir0 >= ir1) return;                                                                                     \
+                                                                                                                \
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                      \
+    dma_queue * dma_q = octx->ctx->dma[ith];                                                                    \
+                                                                                                                \
+    uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS];                                                           \
+    for (uint32_t s = 0; s < n_ranks; s++) {                                                                    \
+        src_spad_base[s] = actx->src_spad_base[s] + (ith * actx->vtcm_size_per_thread);                         \
+    }                                                                                                           \
+    uint8_t * dst_spad_base = actx->dst_spad_base + (ith * actx->vtcm_size_per_thread);                         \
+    uint8_t * res_spad_base = HAS_ADD ? (actx->res_spad_base + (ith * actx->vtcm_size_per_thread)) : NULL;      \
+                                                                                                                \
+    const size_t spad_half = actx->vtcm_size_per_thread / 2;                                                    \
+    uint32_t ir_prefetch = ir0;                                                                                 \
+    int spad_idx = 0;                                                                                           \
+                                                                                                                \
+    for (int k = 0; k < 2 && ir_prefetch < ir1; k++) {                                                          \
+        uint32_t cur_elems = MIN(block_elems, ir1 - ir_prefetch);                                               \
+        size_t   cur_bytes = cur_elems * sizeof(TYPE);                                                          \
+        uint8_t * d_spad = dst_spad_base + spad_idx * spad_half;                                                \
+        for (uint32_t d = 0; d < n_dsts; d++) {                                                                 \
+            dma_addr_t d_ddr = octx->dsts[d]->data + ir_prefetch * sizeof(TYPE);                                \
+            dma_queue_push(dma_q, dma_make_data(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 0);            \
+        }                                                                                                       \
+        for (uint32_t s = 0; s < n_ranks; s++) {                                                                \
+            uint8_t * s_spad = src_spad_base[s] + spad_idx * spad_half;                                         \
+            const dma_addr_t s_ddr = octx->src[s]->data + ir_prefetch * sizeof(TYPE);                           \
+            dma_queue_push(dma_q, dma_make_data(s_spad, s_ddr), cur_bytes, cur_bytes, cur_bytes, 1);            \
+        }                                                                                                       \
+        if (HAS_ADD) {                                                                                          \
+            uint8_t * r_spad = res_spad_base + spad_idx * spad_half;                                            \
+            const dma_addr_t r_ddr = octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE);                 \
+            dma_queue_push(dma_q, dma_make_data(r_spad, r_ddr), cur_bytes, cur_bytes, cur_bytes, 1);            \
+        }                                                                                                       \
+        ir_prefetch += cur_elems;                                                                               \
+        spad_idx ^= 1;                                                                                          \
+    }                                                                                                           \
+                                                                                                                \
+    for (uint32_t ir = ir0; ir < ir1; ) {                                                                       \
+        uint32_t cur_elems = MIN(block_elems, ir1 - ir);                                                        \
+        size_t   cur_bytes = cur_elems * sizeof(TYPE);                                                          \
+        uint8_t * d_spad = NULL;                                                                                \
+        for (uint32_t d = 0; d < n_dsts; d++) {                                                                 \
+            d_spad = (uint8_t *) dma_queue_pop(dma_q).src;                                                      \
+        }                                                                                                       \
+        uint8_t * s_spad[HTP_ALLREDUCE_MAX_RANKS];                                                              \
+        for (uint32_t s = 0; s < n_ranks; s++) {                                                                \
+            s_spad[s] = (uint8_t *) dma_queue_pop(dma_q).dst;                                                   \
+        }                                                                                                       \
+        uint8_t * r_spad = HAS_ADD ? (uint8_t *) dma_queue_pop(dma_q).dst : NULL;                               \
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);                                       \
+        HVX_ADD_FN(d_spad, s_spad[0], s_spad[1], cur_elems);                                                    \
+        for (uint32_t s = 2; s < n_ranks; s++) {                                                                \
+            HVX_ADD_FN(d_spad, d_spad, s_spad[s], cur_elems);                                                   \
+        }                                                                                                       \
+        if (HAS_ADD) {                                                                                          \
+            HVX_ADD_FN(d_spad, d_spad, r_spad, cur_elems);                                                      \
+        }                                                                                                       \
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir);                                        \
+        for (uint32_t d = 0; d < n_dsts; d++) {                                                                 \
+            dma_addr_t d_ddr = octx->dsts[d]->data + ir * sizeof(TYPE);                                         \
+            dma_queue_push(dma_q, dma_make_data(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 1);            \
+        }                                                                                                       \
+        if (ir_prefetch < ir1) {                                                                                \
+            uint32_t next_elems = MIN(block_elems, ir1 - ir_prefetch);                                          \
+            size_t   next_bytes = next_elems * sizeof(TYPE);                                                    \
+            for (uint32_t s = 0; s < n_ranks; s++) {                                                            \
+                const dma_addr_t s_next = octx->src[s]->data + ir_prefetch * sizeof(TYPE);                      \
+                dma_queue_push(dma_q, dma_make_data(s_spad[s], s_next), next_bytes, next_bytes, next_bytes, 1); \
+            }                                                                                                   \
+            if (HAS_ADD) {                                                                                      \
+                const dma_addr_t r_next = octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE);            \
+                dma_queue_push(dma_q, dma_make_data(r_spad, r_next), next_bytes, next_bytes, next_bytes, 1);    \
+            }                                                                                                   \
+            ir_prefetch += next_elems;                                                                          \
+        }                                                                                                       \
+        ir += cur_elems;                                                                                        \
+    }                                                                                                           \
+    dma_queue_flush(dma_q);                                                                                     \
 }

 DEFINE_ALLREDUCE_THREAD_DMA_1D(f16,     __fp16, hvx_add_f16_aaa, 0)
@@ -162,9 +156,12 @@ static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith,
    struct htp_thread_trace * tr = &octx->ctx->trace[ith];                                                                                        \
    dma_queue * dma_q = octx->ctx->dma[ith];                                                                                                      \
                                                                                                                                                  \
-    const size_t vtcm_thread_offset = ith * actx->vtcm_size_per_thread;                                                                           \
-    uint8_t * dst_spad_base = actx->dst_spad_base + vtcm_thread_offset;                                                                           \
-    uint8_t * res_spad_base = HAS_ADD ? (IS_ROW_BCAST ? actx->res_spad_base : (actx->res_spad_base + vtcm_thread_offset)) : NULL;                 \
+    uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS];                                                                                             \
+    for (uint32_t s = 0; s < n_ranks; s++) {                                                                                                      \
+        src_spad_base[s] = actx->src_spad_base[s] + (ith * actx->vtcm_size_per_thread);                                                           \
+    }                                                                                                                                             \
+    uint8_t * dst_spad_base = actx->dst_spad_base + (ith * actx->vtcm_size_per_thread);                                                           \
+    uint8_t * res_spad_base = HAS_ADD ? (IS_ROW_BCAST ? actx->res_spad_base : (actx->res_spad_base + (ith * actx->vtcm_size_per_thread))) : NULL; \
                                                                                                                                                  \
    const size_t spad_half = actx->vtcm_size_per_thread / 2;                                                                                      \
    uint32_t r_prefetch = r0;                                                                                                                     \
@@ -178,7 +175,7 @@ static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith,
            dma_queue_push(dma_q, dma_make_data(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, 0);                            \
        }                                                                                                                                         \
        for (uint32_t s = 0; s < n_ranks; s++) {                                                                                                  \
-            uint8_t * s_spad = actx->src_spad_base[s] + vtcm_thread_offset + spad_idx * spad_half;                                                \
+            uint8_t * s_spad = src_spad_base[s] + spad_idx * spad_half;                                                                           \
            const dma_addr_t s_ddr = octx->src[s]->data + r_prefetch * octx->src[s]->nb[1];                                                       \
            dma_queue_push(dma_q, dma_make_data(s_spad, s_ddr), row_size_aligned, octx->src[s]->nb[1], row_bytes, cur_rows);                      \
        }                                                                                                                                         \
@@ -191,30 +188,25 @@ static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith,
        spad_idx ^= 1;                                                                                                                            \
    }                                                                                                                                             \
                                                                                                                                                  \
-    int comp_spad_idx = 0;                                                                                                                        \
    for (uint32_t r = r0; r < r1; ) {                                                                                                             \
        uint32_t cur_rows = MIN(block_rows, r1 - r);                                                                                              \
+        uint8_t * d_spad = NULL;                                                                                                                  \
        for (uint32_t d = 0; d < n_dsts; d++) {                                                                                                   \
-            dma_queue_pop(dma_q);                                                                                                                 \
+            d_spad = (uint8_t *) dma_queue_pop(dma_q).src;                                                                                        \
        }                                                                                                                                         \
+        uint8_t * s_spad[HTP_ALLREDUCE_MAX_RANKS];                                                                                                \
        for (uint32_t s = 0; s < n_ranks; s++) {                                                                                                  \
-            dma_queue_pop(dma_q);                                                                                                                 \
+            s_spad[s] = (uint8_t *) dma_queue_pop(dma_q).dst;                                                                                     \
        }                                                                                                                                         \
-        if (HAS_ADD && !IS_ROW_BCAST) {                                                                                                           \
-            dma_queue_pop(dma_q);                                                                                                                 \
-        }                                                                                                                                         \
-        uint8_t * d_spad  = dst_spad_base + comp_spad_idx * spad_half;                                                                            \
-        uint8_t * s0_spad = actx->src_spad_base[0] + vtcm_thread_offset + comp_spad_idx * spad_half;                                              \
-        uint8_t * s1_spad = actx->src_spad_base[1] + vtcm_thread_offset + comp_spad_idx * spad_half;                                              \
-        uint8_t * r_spad  = (HAS_ADD && !IS_ROW_BCAST) ? (res_spad_base + comp_spad_idx * spad_half) : NULL;                                      \
+        uint8_t * r_spad = (HAS_ADD && !IS_ROW_BCAST) ? (uint8_t *) dma_queue_pop(dma_q).dst : NULL;                                              \
        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r);                                                                          \
        for (uint32_t row = 0; row < cur_rows; row++) {                                                                                           \
            uint8_t * d_row = d_spad + row * row_size_aligned;                                                                                    \
-            const uint8_t * s0_row = s0_spad + row * row_size_aligned;                                                                            \
-            const uint8_t * s1_row = s1_spad + row * row_size_aligned;                                                                            \
+            const uint8_t * s0_row = s_spad[0] + row * row_size_aligned;                                                                          \
+            const uint8_t * s1_row = s_spad[1] + row * row_size_aligned;                                                                          \
            HVX_ADD_FN(d_row, s0_row, s1_row, ne0);                                                                                               \
            for (uint32_t s = 2; s < n_ranks; s++) {                                                                                              \
-                const uint8_t * ss_row = actx->src_spad_base[s] + vtcm_thread_offset + comp_spad_idx * spad_half + row * row_size_aligned;        \
+                const uint8_t * ss_row = s_spad[s] + row * row_size_aligned;                                                                      \
                HVX_ADD_FN(d_row, d_row, ss_row, ne0);                                                                                            \
            }                                                                                                                                     \
            if (HAS_ADD) {                                                                                                                        \
@@ -230,18 +222,15 @@ static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith,
        if (r_prefetch < r1) {                                                                                                                    \
            uint32_t next_rows = MIN(block_rows, r1 - r_prefetch);                                                                                \
            for (uint32_t s = 0; s < n_ranks; s++) {                                                                                              \
-                uint8_t * s_spad = actx->src_spad_base[s] + vtcm_thread_offset + comp_spad_idx * spad_half;                                       \
                const dma_addr_t s_next = octx->src[s]->data + r_prefetch * octx->src[s]->nb[1];                                                  \
-                dma_queue_push(dma_q, dma_make_data(s_spad, s_next), row_size_aligned, octx->src[s]->nb[1], row_bytes, next_rows);                \
+                dma_queue_push(dma_q, dma_make_data(s_spad[s], s_next), row_size_aligned, octx->src[s]->nb[1], row_bytes, next_rows);             \
            }                                                                                                                                     \
            if (HAS_ADD && !IS_ROW_BCAST) {                                                                                                       \
-                uint8_t * r_spad_next = res_spad_base + comp_spad_idx * spad_half;                                                                \
                const dma_addr_t r_next = octx->src[2 * n_ranks]->data + r_prefetch * octx->src[2 * n_ranks]->nb[1];                              \
-                dma_queue_push(dma_q, dma_make_data(r_spad_next, r_next), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, next_rows); \
+                dma_queue_push(dma_q, dma_make_data(r_spad, r_next), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, next_rows);      \
            }                                                                                                                                     \
            r_prefetch += next_rows;                                                                                                              \
        }                                                                                                                                         \
-        comp_spad_idx ^= 1;                                                                                                                       \
        r += cur_rows;                                                                                                                            \
    }                                                                                                                                             \
    dma_queue_flush(dma_q);                                                                                                                       \
@@ -259,10 +248,9 @@ static int validate_allreduce(
    const struct htp_allreduce_kernel_params * kparams,
    uint32_t n_ranks
 ) {
-    if (kparams->n_threads == 0 || (uint32_t) kparams->n_threads > octx->ctx->n_threads) {
+    if (!htp_ops_context_set_n_threads(octx, (uint32_t) kparams->n_threads)) {
        return HTP_STATUS_INVAL_PARAMS;
    }
-    octx->n_threads = (uint32_t) kparams->n_threads;

    if (kparams->vtcm_size_per_thread <= 0 || kparams->vtcm_size <= 0) {
        return HTP_STATUS_INVAL_PARAMS;
@@ -318,7 +306,6 @@ int op_allreduce(struct htp_ops_context * octx) {

    const bool has_add = (octx->op == HTP_OP_ALLREDUCE_ADD);
    const uint32_t nelem = dst->ne[0] * dst->ne[1] * dst->ne[2] * dst->ne[3];
-    const int32_t  mode  = kparams->mode;

    // 1. Entry Barrier: Synchronize all ranks before reading
    struct htp_thread_trace * tr0 = &octx->ctx->trace[0];
@@ -432,11 +419,6 @@ int op_allreduce(struct htp_ops_context * octx) {
    // 4. Exit Barrier: Synchronize all ranks after writing
    htp_trace_event_start(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit);

-    // drain fan-out DMA writes before exit barrier
-    if (mode == HTP_ALLREDUCE_SHARDED_FANOUT) {
-        asm volatile ("syncht" : : : "memory");
-    }
-
    htp_fence_write(my_fence, fence_seq_exit, octx->status);

    for (uint32_t j = 0; j < n_ranks; j++) {
@@ -17,11 +17,6 @@ enum htp_allreduce_kernel_type {
    HTP_ALLREDUCE_KERNEL_DMA_2D,
 };

-enum htp_allreduce_mode {
-    HTP_ALLREDUCE_FULL           = 0, // all ranks reduce full tensor
-    HTP_ALLREDUCE_SHARDED_FANOUT = 1, // each rank reduces 1/N shard and fans out to all buffers
-};
-
 static inline size_t htp_allreduce_vtcm_buffer_count(
    uint32_t n_ranks,
    uint32_t n_threads,
@@ -47,7 +42,6 @@ struct htp_allreduce_kernel_params {
    int32_t rank_nelem;
    int32_t n_dsts;
    int32_t is_row_bcast;
-    int32_t mode;                 // enum htp_allreduce_mode: FULL or SHARDED_FANOUT
 };

 #ifdef __cplusplus
@@ -26,12 +26,6 @@ struct htp_concat_context {
    struct fastdiv_values div_ne2;
 };

-static inline dma_addr_t concat_plane_addr(const struct htp_tensor * t, uint32_t p, const struct fastdiv_values * div_ne2, uint32_t ne2) {
-    const uint32_t i3 = fastdiv(p, div_ne2);
-    const uint32_t i2 = p - i3 * ne2;
-    return t->data + i2 * t->nb[2] + i3 * t->nb[3];
-}
-
 static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void * data) {
    struct htp_concat_context * cctx = (struct htp_concat_context *) data;
    struct htp_ops_context * octx = cctx->octx;
@@ -56,74 +50,58 @@ static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void *
    const uint32_t block_i = 32;
    const uint32_t spad1_stride = block_i * sizeof(float);

-    const HVX_Vector offsets = hvx_vec_gather_offsets_w(spad1_stride);
+    int32_t offsets[32] __attribute__((aligned(128)));
+    for(int k=0; k<32; k++) {
+        offsets[k] = k * spad1_stride;
+    }
+    HVX_Vector vv = *(HVX_Vector*)offsets;
    const uint32_t src1_ne0_padded = hex_round_up(src1_ne0, 32);
-    const uint32_t src0_row_bytes  = src0_ne0 * sizeof(float);
-    const uint32_t src0_row_padded = hex_round_up(src0_row_bytes, VLEN);
-    const uint32_t src0_pre        = src0_row_padded - src0_row_bytes;
-    const uint32_t spad0_row_bytes = src0_row_padded + src1_ne0_padded * sizeof(float);
+    const uint32_t spad0_row_bytes = hex_round_up((src0_ne0 + src1_ne0_padded) * sizeof(float), VLEN);
+    uint32_t mu = src1_ne0_padded * spad1_stride;

    struct htp_thread_trace * tr = &octx->ctx->trace[ith];

-    const struct fastdiv_values * div_ne2 = &cctx->div_ne2;
-    const uint32_t ne2 = dst->ne[2];
+    for (uint32_t p = 0; p < cctx->nplanes; p++) {
+        const uint32_t i3 = p / dst->ne[2];
+        const uint32_t i2 = p - i3 * dst->ne[2];
+        const dma_addr_t src0_plane = src0->data + i2 * src0->nb[2] + i3 * src0->nb[3];
+        const dma_addr_t src1_plane = src1->data + i2 * src1->nb[2] + i3 * src1->nb[3];
+        const dma_addr_t dst_plane  = dst->data  + i2 * dst->nb[2]  + i3 * dst->nb[3];

-    uint32_t p = 0;
-    uint32_t i = start_i;
+        for (uint32_t i = start_i; i < end_i; i += block_i) {
+            uint32_t current_block_i = (end_i - i < block_i) ? (end_i - i) : block_i;

-    const dma_addr_t src1_addr = concat_plane_addr(src1, p, div_ne2, ne2) + i * src1->nb[1];
-    dma_queue_push(dma_q, dma_make_data(spad1_base, src1_addr), spad1_stride, src1->nb[0], MIN(end_i - i, block_i) * sizeof(float), src1_ne0);
+            uint32_t src1_width_bytes = current_block_i * sizeof(float);
+            const dma_addr_t src1_addr = src1_plane + i * src1->nb[1];
+            dma_queue_push(dma_q, dma_make_data(spad1_base, src1_addr), spad1_stride, src1->nb[0], src1_width_bytes, src1_ne0);

-    const dma_addr_t src0_addr = concat_plane_addr(src0, p, div_ne2, ne2) + i * src0->nb[1];
-    dma_queue_push(dma_q, dma_make_data(spad0_base + src0_pre, src0_addr), spad0_row_bytes, src0->nb[1], src0_row_bytes, MIN(end_i - i, block_i));
+            uint32_t src0_row_bytes = src0_ne0 * sizeof(float);
+            const dma_addr_t src0_addr = src0_plane + i * src0->nb[1];
+            dma_queue_push(dma_q, dma_make_data(spad0_base, src0_addr), spad0_row_bytes, src0->nb[1], src0_row_bytes, current_block_i);

-    dma_queue_pop(dma_q); // src1
+            dma_queue_pop(dma_q); // src1

-    while (p < cctx->nplanes) {
-        const uint32_t current_block_i = MIN(end_i - i, block_i);
+            HVX_Vector * vtcm_tmp = (HVX_Vector *)(spad1_base + src1_ne0_padded * spad1_stride);

-        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
-        for (uint32_t j = 0; j < src1_ne0; j += 32) {
-            const uint8_t * src_ptr = spad1_base + j * spad1_stride;
-            uint8_t * dst_ptr = spad0_base + src0_row_padded + j * sizeof(float);
-            hvx_transpose_32x32_w_gather(dst_ptr, spad0_row_bytes, src_ptr, spad1_stride, offsets, current_block_i, MIN(src1_ne0 - j, 32));
+            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
+            for (uint32_t j = 0; j < src1_ne0_padded; j += 32) {
+                #pragma unroll(4)
+                for (uint32_t ii = 0; ii < current_block_i; ii++) {
+                    size_t rt = (size_t)(spad1_base + j * spad1_stride + ii * sizeof(float));
+                    Q6_vgather_ARMVw(&vtcm_tmp[ii], rt, mu, vv);
+                    uint8_t * dst_ptr = spad0_base + ii * spad0_row_bytes + (src0_ne0 + j) * sizeof(float);
+                    hvx_vmemu(dst_ptr) = vtcm_tmp[ii];
+                }
+            }
+            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
+
+            dma_queue_pop(dma_q); // src0
+
+            const dma_addr_t dst_addr = dst_plane + i * dst->nb[1];
+            dma_queue_push(dma_q, dma_make_data(dst_addr, spad0_base), dst->nb[1], spad0_row_bytes, (src0_ne0 + src1_ne0) * sizeof(float), current_block_i);
+
+            dma_queue_pop(dma_q);
        }
-        hvx_gather_sync(spad0_base + src0_row_padded);
-        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
-
-        uint32_t np = p;
-        uint32_t ni = i + block_i;
-        if (ni >= end_i) {
-            ni = start_i;
-            np++;
-        }
-        const bool has_next = np < cctx->nplanes;
-        const uint32_t next_block_i = MIN(end_i - ni, block_i);
-
-        // spad1 is free after the gather sync, prefetch next src1 ahead of the dst write
-        if (has_next) {
-            const dma_addr_t nsrc1_addr = concat_plane_addr(src1, np, div_ne2, ne2) + ni * src1->nb[1];
-            dma_queue_push(dma_q, dma_make_data(spad1_base, nsrc1_addr), spad1_stride, src1->nb[0], next_block_i * sizeof(float), src1_ne0);
-        }
-
-        dma_queue_pop(dma_q); // src0
-
-        const dma_addr_t dst_addr = concat_plane_addr(dst, p, div_ne2, ne2) + i * dst->nb[1];
-        dma_queue_push(dma_q, dma_make_data(dst_addr, spad0_base + src0_pre), dst->nb[1], spad0_row_bytes, (src0_ne0 + src1_ne0) * sizeof(float), current_block_i);
-
-        if (has_next) {
-            dma_queue_pop(dma_q); // next src1
-        }
-        dma_queue_pop(dma_q); // dst
-
-        // spad0 is free after the dst write
-        if (has_next) {
-            const dma_addr_t nsrc0_addr = concat_plane_addr(src0, np, div_ne2, ne2) + ni * src0->nb[1];
-            dma_queue_push(dma_q, dma_make_data(spad0_base + src0_pre, nsrc0_addr), spad0_row_bytes, src0->nb[1], src0_row_bytes, next_block_i);
-        }
-
-        p = np;
-        i = ni;
    }
 }

@@ -151,74 +129,58 @@ static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void *
    const uint32_t block_i = 64;
    const uint32_t spad1_stride = block_i * sizeof(__fp16);

-    const HVX_Vector offsets = hvx_vec_gather_offsets_h(spad1_stride);
+    int16_t offsets[64] __attribute__((aligned(128)));
+    for(int k=0; k<64; k++) {
+        offsets[k] = k * spad1_stride;
+    }
+    HVX_Vector vv = *(HVX_Vector*)offsets;
    const uint32_t src1_ne0_padded = hex_round_up(src1_ne0, 64);
-    const uint32_t src0_row_bytes  = src0_ne0 * sizeof(__fp16);
-    const uint32_t src0_row_padded = hex_round_up(src0_row_bytes, VLEN);
-    const uint32_t src0_pre        = src0_row_padded - src0_row_bytes;
-    const uint32_t spad0_row_bytes = src0_row_padded + src1_ne0_padded * sizeof(__fp16);
+    const uint32_t spad0_row_bytes = hex_round_up((src0_ne0 + src1_ne0_padded) * sizeof(__fp16), VLEN);
+    uint32_t mu = src1_ne0_padded * spad1_stride;

    struct htp_thread_trace * tr = &octx->ctx->trace[ith];

-    const struct fastdiv_values * div_ne2 = &cctx->div_ne2;
-    const uint32_t ne2 = dst->ne[2];
+    for (uint32_t p = 0; p < cctx->nplanes; p++) {
+        const uint32_t i3 = p / dst->ne[2];
+        const uint32_t i2 = p - i3 * dst->ne[2];
+        const dma_addr_t src0_plane = src0->data + i2 * src0->nb[2] + i3 * src0->nb[3];
+        const dma_addr_t src1_plane = src1->data + i2 * src1->nb[2] + i3 * src1->nb[3];
+        const dma_addr_t dst_plane  = dst->data  + i2 * dst->nb[2]  + i3 * dst->nb[3];

-    uint32_t p = 0;
-    uint32_t i = start_i;
+        for (uint32_t i = start_i; i < end_i; i += block_i) {
+            uint32_t current_block_i = (end_i - i < block_i) ? (end_i - i) : block_i;

-    const dma_addr_t src1_addr = concat_plane_addr(src1, p, div_ne2, ne2) + i * src1->nb[1];
-    dma_queue_push(dma_q, dma_make_data(spad1_base, src1_addr), spad1_stride, src1->nb[0], MIN(end_i - i, block_i) * sizeof(__fp16), src1_ne0);
+            uint32_t src1_width_bytes = current_block_i * sizeof(__fp16);
+            const dma_addr_t src1_addr = src1_plane + i * src1->nb[1];
+            dma_queue_push(dma_q, dma_make_data(spad1_base, src1_addr), spad1_stride, src1->nb[0], src1_width_bytes, src1_ne0);

-    const dma_addr_t src0_addr = concat_plane_addr(src0, p, div_ne2, ne2) + i * src0->nb[1];
-    dma_queue_push(dma_q, dma_make_data(spad0_base + src0_pre, src0_addr), spad0_row_bytes, src0->nb[1], src0_row_bytes, MIN(end_i - i, block_i));
+            uint32_t src0_row_bytes = src0_ne0 * sizeof(__fp16);
+            const dma_addr_t src0_addr = src0_plane + i * src0->nb[1];
+            dma_queue_push(dma_q, dma_make_data(spad0_base, src0_addr), spad0_row_bytes, src0->nb[1], src0_row_bytes, current_block_i);

-    dma_queue_pop(dma_q); // src1
+            dma_queue_pop(dma_q); // src1

-    while (p < cctx->nplanes) {
-        const uint32_t current_block_i = MIN(end_i - i, block_i);
+            HVX_Vector * vtcm_tmp = (HVX_Vector *)(spad1_base + src1_ne0_padded * spad1_stride);

-        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
-        for (uint32_t j = 0; j < src1_ne0; j += 64) {
-            const uint8_t * src_ptr = spad1_base + j * spad1_stride;
-            uint8_t * dst_ptr = spad0_base + src0_row_padded + j * sizeof(__fp16);
-            hvx_transpose_64x64_h_gather(dst_ptr, spad0_row_bytes, src_ptr, spad1_stride, offsets, current_block_i, MIN(src1_ne0 - j, 64));
+            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
+            for (uint32_t j = 0; j < src1_ne0_padded; j += 64) {
+                #pragma unroll(4)
+                for (uint32_t ii = 0; ii < current_block_i; ii++) {
+                    size_t rt = (size_t)(spad1_base + j * spad1_stride + ii * sizeof(__fp16));
+                    Q6_vgather_ARMVh(&vtcm_tmp[ii], rt, mu, vv);
+                    uint8_t * dst_ptr = spad0_base + ii * spad0_row_bytes + (src0_ne0 + j) * sizeof(__fp16);
+                    hvx_vmemu(dst_ptr) = vtcm_tmp[ii];
+                }
+            }
+            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
+
+            dma_queue_pop(dma_q); // src0
+
+            const dma_addr_t dst_addr = dst_plane + i * dst->nb[1];
+            dma_queue_push(dma_q, dma_make_data(dst_addr, spad0_base), dst->nb[1], spad0_row_bytes, (src0_ne0 + src1_ne0) * sizeof(__fp16), current_block_i);
+
+            dma_queue_pop(dma_q);
        }
-        hvx_gather_sync(spad0_base + src0_row_padded);
-        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
-
-        uint32_t np = p;
-        uint32_t ni = i + block_i;
-        if (ni >= end_i) {
-            ni = start_i;
-            np++;
-        }
-        const bool has_next = np < cctx->nplanes;
-        const uint32_t next_block_i = MIN(end_i - ni, block_i);
-
-        // spad1 is free after the gather sync, prefetch next src1 ahead of the dst write
-        if (has_next) {
-            const dma_addr_t nsrc1_addr = concat_plane_addr(src1, np, div_ne2, ne2) + ni * src1->nb[1];
-            dma_queue_push(dma_q, dma_make_data(spad1_base, nsrc1_addr), spad1_stride, src1->nb[0], next_block_i * sizeof(__fp16), src1_ne0);
-        }
-
-        dma_queue_pop(dma_q); // src0
-
-        const dma_addr_t dst_addr = concat_plane_addr(dst, p, div_ne2, ne2) + i * dst->nb[1];
-        dma_queue_push(dma_q, dma_make_data(dst_addr, spad0_base + src0_pre), dst->nb[1], spad0_row_bytes, (src0_ne0 + src1_ne0) * sizeof(__fp16), current_block_i);
-
-        if (has_next) {
-            dma_queue_pop(dma_q); // next src1
-        }
-        dma_queue_pop(dma_q); // dst
-
-        // spad0 is free after the dst write
-        if (has_next) {
-            const dma_addr_t nsrc0_addr = concat_plane_addr(src0, np, div_ne2, ne2) + ni * src0->nb[1];
-            dma_queue_push(dma_q, dma_make_data(spad0_base + src0_pre, nsrc0_addr), spad0_row_bytes, src0->nb[1], src0_row_bytes, next_block_i);
-        }
-
-        p = np;
-        i = ni;
    }
 }

@@ -238,7 +200,7 @@ static void concat_generic(unsigned int nth, unsigned int ith, void * data) {
    // Per-device element range aligned to prevent false sharing
    const uint32_t elem_start = cctx->elem_start;
    const uint32_t nelems     = cctx->nelems;
-    const uint32_t chunk_size = fastdiv(nelems + nth - 1, &octx->n_threads_div);
+    const uint32_t chunk_size = (nelems + nth - 1) / nth;

    const uint32_t start_idx = MIN(elem_start + ith * chunk_size, elem_start + nelems);
    const uint32_t end_idx   = MIN(start_idx + chunk_size, elem_start + nelems);
@@ -396,11 +358,10 @@ int op_concat(struct htp_ops_context * octx) {
        uint32_t spad1_stride = block_i * type_size;

        uint32_t src1_ne0_padded = hex_round_up(src1->ne[0], block_i);
-        // src0 row is right-aligned to VLEN so the gathered src1 part starts aligned
-        uint32_t spad0_row_bytes = hex_round_up(src0->ne[0] * type_size, VLEN) + src1_ne0_padded * type_size;
+        uint32_t spad0_row_bytes = hex_round_up((src0->ne[0] + src1_ne0_padded) * type_size, VLEN);

        octx->src0_spad.size_per_thread = block_i * spad0_row_bytes;
-        octx->src1_spad.size_per_thread = src1_ne0_padded * spad1_stride;
+        octx->src1_spad.size_per_thread = src1_ne0_padded * spad1_stride + block_i * VLEN;

        octx->src0_spad.size = n_threads * octx->src0_spad.size_per_thread;
        octx->src1_spad.size = n_threads * octx->src1_spad.size_per_thread;
--- a/Show More
+++ b/Show More