graph: store the conv state of every rollback arch in one set_rows

The rollback slots of the causal conv history were written with one
copy per slot in qwen4exp, glm5-next, kimi-k3, bailingmoe3, lfm and
mamba2, while delta_net already wrote them with a single set_rows.
build_conv_state_store now does this for all of them: one copy without
rollback, otherwise one set_rows of the overlapping windows into the
min(n, n_rs_seq + 1) slots a rollback can read.

The slot order {n_rs_seq, ..., 0} is a constant of the recurrent
memory, kept in its own buffer next to the state, and each graph takes
a view of its tail. Its size no longer depends on the ubatch, so the
reserved graph covers every ubatch shape, and the graph input that
held it is gone.

A model that is not causal never rolls back, so the context clamps
n_rs_seq to 0 for it, like for an arch without rollback support.

.github/workflows/behavior-check.yml

@@ -13,6 +13,8 @@ on:
       '**/*.hpp',
       '**/*.c',
       '**/*.cpp',
+      '**/*.py',
+      '**/requirements*.txt',
     ]
 
   pull_request:
@@ -24,7 +26,9 @@ on:
       '**/*.h',
       '**/*.hpp',
       '**/*.c',
-      '**/*.cpp'
+      '**/*.cpp',
+      '**/*.py',
+      '**/requirements*.txt',
     ]
 
 cache-mode: none
@@ -69,7 +73,7 @@ jobs:
         run: |
           python3 -m venv .venv
           source .venv/bin/activate
-          pip install jinja2==3.1.6
+          pip install -r requirements/requirements-convert_hf_to_gguf.txt
 
       - name: ccache-buckets-restore
         uses: ./.github/actions/ccache-buckets
@@ -134,7 +138,7 @@ jobs:
         run: |
           python3 -m venv .venv
           source .venv/Scripts/activate
-          pip install jinja2==3.1.6
+          pip install -r requirements/requirements-convert_hf_to_gguf.txt
 
       - name: Build
         id: cmake_build
@@ -156,3 +160,13 @@ jobs:
           source .venv/Scripts/activate
           cd build
           ctest -C Release -L behavior -LE 'model|backend' --verbose --timeout 900
+
+      - name: ccache-clear
+        uses: ./.github/actions/ccache-clear
+        env:
+          GH_TOKEN: ${{ github.token }}
+        with:
+          key: behavior-check-windows-2025
+          older: 5m
+          min: 1
+          dry-run: ${{ github.event_name != 'push' || github.ref != 'refs/heads/master' }}

.github/workflows/build-apple-daily.yml

@@ -0,0 +1,160 @@
+name: CI (apple - daily)
+
+on:
+  workflow_dispatch: # allows manual triggering
+  schedule:
+    - cron: '0 0 * * *'
+
+cache-mode: none
+permissions:
+  contents: read
+
+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:
+  macos-latest-ios-xcode:
+    runs-on: macos-latest
+
+    permissions:
+      actions: write
+      contents: read
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v6
+
+      - name: Setup Xcode
+        uses: ggml-org/setup-xcode@v1
+        with:
+          xcode-version: latest-stable
+
+      - name: Build
+        id: cmake_build
+        run: |
+          sysctl -a
+          cmake -B build -G Xcode \
+            -DGGML_METAL_EMBED_LIBRARY=ON \
+            -DLLAMA_OPENSSL=OFF \
+            -DLLAMA_BUILD_APP=OFF \
+            -DLLAMA_BUILD_EXAMPLES=OFF \
+            -DLLAMA_BUILD_TOOLS=OFF \
+            -DLLAMA_BUILD_TESTS=OFF \
+            -DLLAMA_BUILD_SERVER=OFF \
+            -DCMAKE_SYSTEM_NAME=iOS \
+            -DCMAKE_OSX_DEPLOYMENT_TARGET=14.0 \
+            -DCMAKE_XCODE_ATTRIBUTE_DEVELOPMENT_TEAM=ggml
+          cmake --build build --config Release -j $(sysctl -n hw.logicalcpu) -- CODE_SIGNING_ALLOWED=NO
+
+      - name: xcodebuild for swift package
+        id: xcodebuild
+        run: |
+          ./build-xcframework.sh
+
+      - name: Upload xcframework artifact
+        uses: actions/upload-artifact@v6
+        with:
+          name: llama-xcframework
+          path: build-apple/llama.xcframework/
+          retention-days: 1
+
+      - name: Build Xcode project
+        run: |
+          xcodebuild -downloadPlatform iOS
+          xcodebuild -project examples/llama.swiftui/llama.swiftui.xcodeproj -scheme llama.swiftui -sdk iphoneos CODE_SIGNING_REQUIRED=NO CODE_SIGN_IDENTITY= -destination 'generic/platform=iOS' FRAMEWORK_FOLDER_PATH=./build-ios build
+
+  macos-latest-tvos:
+    runs-on: macos-latest
+
+    steps:
+      - name: Clone
+        id: checkout
+        uses: actions/checkout@v6
+
+      - name: Build
+        id: cmake_build
+        run: |
+          sysctl -a
+          cmake -B build -G Xcode \
+            -DGGML_METAL_EMBED_LIBRARY=ON \
+            -DLLAMA_BUILD_COMMON=OFF \
+            -DLLAMA_BUILD_APP=OFF \
+            -DLLAMA_BUILD_EXAMPLES=OFF \
+            -DLLAMA_BUILD_TOOLS=OFF \
+            -DLLAMA_BUILD_TESTS=OFF \
+            -DLLAMA_BUILD_SERVER=OFF \
+            -DCMAKE_SYSTEM_NAME=tvOS \
+            -DCMAKE_OSX_DEPLOYMENT_TARGET=14.0 \
+            -DCMAKE_XCODE_ATTRIBUTE_DEVELOPMENT_TEAM=ggml
+          cmake --build build --config Release -j $(sysctl -n hw.logicalcpu) -- CODE_SIGNING_ALLOWED=NO
+
+  macos-latest-visionos:
+    runs-on: macos-latest
+
+    steps:
+      - name: Clone
+        id: checkout
+        uses: actions/checkout@v6
+
+      - name: Build
+        id: cmake_build
+        run: |
+          sysctl -a
+          cmake -B build -G Xcode \
+            -DGGML_METAL_EMBED_LIBRARY=ON \
+            -DLLAMA_BUILD_COMMON=OFF \
+            -DLLAMA_BUILD_APP=OFF \
+            -DLLAMA_BUILD_EXAMPLES=OFF \
+            -DLLAMA_BUILD_TOOLS=OFF \
+            -DLLAMA_BUILD_TESTS=OFF \
+            -DLLAMA_BUILD_SERVER=OFF \
+            -DCMAKE_SYSTEM_NAME=visionOS \
+            -DCMAKE_OSX_DEPLOYMENT_TARGET=1.0 \
+            -DCMAKE_XCODE_ATTRIBUTE_DEVELOPMENT_TEAM=ggml
+          cmake --build build --config Release -j $(sysctl -n hw.logicalcpu) -- CODE_SIGNING_ALLOWED=NO
+
+  macos-latest-swift:
+    runs-on: macos-latest
+    needs: macos-latest-ios-xcode
+
+    permissions:
+      actions: read
+      contents: read
+
+    strategy:
+      matrix:
+        destination: ['generic/platform=macOS', 'generic/platform=iOS', 'generic/platform=tvOS']
+
+    steps:
+      - name: Clone
+        id: checkout
+        uses: actions/checkout@v6
+
+      - name: Download xcframework artifact
+        uses: actions/download-artifact@v7
+        with:
+          name: llama-xcframework
+          path: build-apple/llama.xcframework/
+
+      - name: Build llama.cpp with CMake
+        id: cmake_build
+        run: |
+          sysctl -a
+          cmake -B build -G Xcode \
+            -DGGML_METAL_EMBED_LIBRARY=ON \
+            -DLLAMA_OPENSSL=OFF \
+            -DLLAMA_BUILD_APP=OFF \
+            -DLLAMA_BUILD_EXAMPLES=OFF \
+            -DLLAMA_BUILD_TOOLS=OFF \
+            -DLLAMA_BUILD_TESTS=OFF \
+            -DLLAMA_BUILD_SERVER=OFF \
+            -DCMAKE_OSX_ARCHITECTURES="arm64;x86_64"
+          cmake --build build --config Release -j $(sysctl -n hw.logicalcpu)

.github/workflows/build-apple.yml

@@ -107,140 +107,3 @@ jobs:
           # ref: https://github.com/ggml-org/llama.cpp/pull/19802#issuecomment-4013704023
           ctest -L main -E "test-llama-archs|test-llama-context|test-speculative-rejection-sampling" --verbose --timeout 900
 
-  macos-latest-ios-xcode:
-    runs-on: macos-latest
-
-    permissions:
-      actions: write
-      contents: read
-
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v6
-
-      - name: Setup Xcode
-        uses: ggml-org/setup-xcode@v1
-        with:
-          xcode-version: latest-stable
-
-      - name: Build
-        id: cmake_build
-        run: |
-          sysctl -a
-          cmake -B build -G Xcode \
-            -DGGML_METAL_EMBED_LIBRARY=ON \
-            -DLLAMA_OPENSSL=OFF \
-            -DLLAMA_BUILD_APP=OFF \
-            -DLLAMA_BUILD_EXAMPLES=OFF \
-            -DLLAMA_BUILD_TOOLS=OFF \
-            -DLLAMA_BUILD_TESTS=OFF \
-            -DLLAMA_BUILD_SERVER=OFF \
-            -DCMAKE_SYSTEM_NAME=iOS \
-            -DCMAKE_OSX_DEPLOYMENT_TARGET=14.0 \
-            -DCMAKE_XCODE_ATTRIBUTE_DEVELOPMENT_TEAM=ggml
-          cmake --build build --config Release -j $(sysctl -n hw.logicalcpu) -- CODE_SIGNING_ALLOWED=NO
-
-      - name: xcodebuild for swift package
-        id: xcodebuild
-        run: |
-          ./build-xcframework.sh
-
-      - name: Upload xcframework artifact
-        uses: actions/upload-artifact@v6
-        with:
-          name: llama-xcframework
-          path: build-apple/llama.xcframework/
-          retention-days: 1
-
-      - name: Build Xcode project
-        run: |
-          xcodebuild -downloadPlatform iOS
-          xcodebuild -project examples/llama.swiftui/llama.swiftui.xcodeproj -scheme llama.swiftui -sdk iphoneos CODE_SIGNING_REQUIRED=NO CODE_SIGN_IDENTITY= -destination 'generic/platform=iOS' FRAMEWORK_FOLDER_PATH=./build-ios build
-
-  macos-latest-tvos:
-    runs-on: macos-latest
-
-    steps:
-      - name: Clone
-        id: checkout
-        uses: actions/checkout@v6
-
-      - name: Build
-        id: cmake_build
-        run: |
-          sysctl -a
-          cmake -B build -G Xcode \
-            -DGGML_METAL_EMBED_LIBRARY=ON \
-            -DLLAMA_BUILD_COMMON=OFF \
-            -DLLAMA_BUILD_APP=OFF \
-            -DLLAMA_BUILD_EXAMPLES=OFF \
-            -DLLAMA_BUILD_TOOLS=OFF \
-            -DLLAMA_BUILD_TESTS=OFF \
-            -DLLAMA_BUILD_SERVER=OFF \
-            -DCMAKE_SYSTEM_NAME=tvOS \
-            -DCMAKE_OSX_DEPLOYMENT_TARGET=14.0 \
-            -DCMAKE_XCODE_ATTRIBUTE_DEVELOPMENT_TEAM=ggml
-          cmake --build build --config Release -j $(sysctl -n hw.logicalcpu) -- CODE_SIGNING_ALLOWED=NO
-
-  macos-latest-visionos:
-    runs-on: macos-latest
-
-    steps:
-      - name: Clone
-        id: checkout
-        uses: actions/checkout@v6
-
-      - name: Build
-        id: cmake_build
-        run: |
-          sysctl -a
-          cmake -B build -G Xcode \
-            -DGGML_METAL_EMBED_LIBRARY=ON \
-            -DLLAMA_BUILD_COMMON=OFF \
-            -DLLAMA_BUILD_APP=OFF \
-            -DLLAMA_BUILD_EXAMPLES=OFF \
-            -DLLAMA_BUILD_TOOLS=OFF \
-            -DLLAMA_BUILD_TESTS=OFF \
-            -DLLAMA_BUILD_SERVER=OFF \
-            -DCMAKE_SYSTEM_NAME=visionOS \
-            -DCMAKE_OSX_DEPLOYMENT_TARGET=1.0 \
-            -DCMAKE_XCODE_ATTRIBUTE_DEVELOPMENT_TEAM=ggml
-          cmake --build build --config Release -j $(sysctl -n hw.logicalcpu) -- CODE_SIGNING_ALLOWED=NO
-
-  macos-latest-swift:
-    runs-on: macos-latest
-    needs: macos-latest-ios-xcode
-
-    permissions:
-      actions: read
-      contents: read
-
-    strategy:
-      matrix:
-        destination: ['generic/platform=macOS', 'generic/platform=iOS', 'generic/platform=tvOS']
-
-    steps:
-      - name: Clone
-        id: checkout
-        uses: actions/checkout@v6
-
-      - name: Download xcframework artifact
-        uses: actions/download-artifact@v7
-        with:
-          name: llama-xcframework
-          path: build-apple/llama.xcframework/
-
-      - name: Build llama.cpp with CMake
-        id: cmake_build
-        run: |
-          sysctl -a
-          cmake -B build -G Xcode \
-            -DGGML_METAL_EMBED_LIBRARY=ON \
-            -DLLAMA_OPENSSL=OFF \
-            -DLLAMA_BUILD_APP=OFF \
-            -DLLAMA_BUILD_EXAMPLES=OFF \
-            -DLLAMA_BUILD_TOOLS=OFF \
-            -DLLAMA_BUILD_TESTS=OFF \
-            -DLLAMA_BUILD_SERVER=OFF \
-            -DCMAKE_OSX_ARCHITECTURES="arm64;x86_64"
-          cmake --build build --config Release -j $(sysctl -n hw.logicalcpu)

.github/workflows/pr-video-demo-label.yml

@@ -0,0 +1,37 @@
+name: "PR video demo label"
+
+on:
+  pull_request_target:
+    types: [opened, edited, reopened, synchronize]
+
+cache-mode: none
+permissions:
+  pull-requests: write
+  issues: write
+
+jobs:
+  video-demo-label:
+    runs-on: ubuntu-slim
+    steps:
+      - name: Add or remove the video demo label
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          PR_URL: ${{ github.event.pull_request.html_url }}
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          # passed through env, so the body is only matched, never run, and is safe to read from a fork
+          BODY: ${{ github.event.pull_request.body }}
+          HAS_LABEL: ${{ contains(github.event.pull_request.labels.*.name, 'video demo') }}
+        run: |
+          # sections written by demo:pr (see tools/ui/demo/cli.ts): the managed comment
+          # markers, or a preview linking to a demo video uploaded for this PR
+          has_demo=false
+          if { grep -qF '<!-- demo:start -->' <<< "$BODY" && grep -qF '<!-- demo:end -->' <<< "$BODY"; } ||
+            grep -qE "/resolve/videos/pr/$PR_NUMBER/[0-9a-f]{10}(-[0-9]+)?\.mp4" <<< "$BODY"; then
+            has_demo=true
+          fi
+
+          if [ "$has_demo" = true ] && [ "$HAS_LABEL" = false ]; then
+            gh pr edit "$PR_URL" --add-label 'video demo'
+          elif [ "$has_demo" = false ] && [ "$HAS_LABEL" = true ]; then
+            gh pr edit "$PR_URL" --remove-label 'video demo'
+          fi

CMakeLists.txt

@@ -150,6 +150,7 @@ option(LLAMA_LLGUIDANCE "llama-common: include LLGuidance library for structured
 # Required for relocatable CMake package
 include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/build-info.cmake)
 include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/common.cmake)
+include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/win-app.cmake)
 
 if (NOT DEFINED LLAMA_BUILD_NUMBER)
     set(LLAMA_BUILD_NUMBER        ${BUILD_NUMBER})
@@ -269,6 +270,8 @@ if (LLAMA_BUILD_MTMD AND NOT (LLAMA_BUILD_COMMON AND LLAMA_BUILD_TOOLS))
     add_subdirectory(tools/mtmd)
 endif()
 
+llama_apply_app_manifests()
+
 #
 # install
 #

README.md

@@ -13,7 +13,7 @@
 [![Docker](https://img.shields.io/github/actions/workflow/status/ggml-org/llama.cpp/docker.yml?label=Docker)](https://github.com/ggml-org/llama.cpp/actions/workflows/docker.yml)
 [![Winget](https://img.shields.io/github/actions/workflow/status/ggml-org/llama.cpp/winget.yml?label=Winget)](https://github.com/ggml-org/llama.cpp/actions/workflows/winget.yml)
 
-[ggml](https://github.com/ggml-org/ggml) / [ops](https://github.com/ggml-org/llama.cpp/blob/master/docs/ops.md) / [maintainer PRs](https://github.com/ggml-org/llama.cpp/issues?q=is%3Apr%20is%3Aopen%20draft%3AFalse%20(author%3Asanmai%20OR%20author%3Argerganov%20OR%20author%3AKitaitiMakoto%20OR%20author%3Adanbev%20OR%20author%3Aaldehir%20OR%20author%3Amax-krasnyansky%20OR%20author%3ACISC%20OR%20author%3Aggerganov%20OR%20author%3Aam17an%20OR%20author%3Ajhen0409%20OR%20author%3Abartowski1182%20OR%20author%3Anikwen%20OR%20author%3Ahipudding%20OR%20author%3Aravi9%20OR%20author%3AServeurpersoCom%20OR%20author%3Athelittlefireman%20OR%20author%3Apwilkin%20OR%20author%3Areeselevine%20OR%20author%3Angxson%20OR%20author%3Ajeffbolznv%20OR%20author%3Amarty1885%20OR%20author%3A0cc4m%20OR%20author%3ATitaniumtown%20OR%20author%3Aangt%20OR%20author%3AIMbackK%20OR%20author%3Aarthw%20OR%20author%3AJohannesGaessler%20OR%20author%3AORippler%20OR%20author%3Aruixiang63%20OR%20author%3Axctan%20OR%20author%3AmasterFoad%20OR%20author%3Aallozaur%20OR%20author%3AAndreas-Krebbel%20OR%20author%3Ayomaytk%20OR%20author%3Aaendk%20OR%20author%3Awine99%20OR%20author%3Agaugarg-nv%20OR%20author%3Ataronaeo%20OR%20author%3Aforforever73%20OR%20author%3Achaxu01%20OR%20author%3Alhez%20OR%20author%3AAlekseiNikiforovIBM%20OR%20author%3Anetrunnereve%20OR%20author%3Afairydreaming%20OR%20author%3Aynankani)%20sort%3Aupdated-desc) / [dev stats](https://github.com/ggml-org/llama.cpp-dev) / [lib llama API](https://github.com/ggml-org/llama.cpp/issues/9289) / [llama-server REST API](https://github.com/ggml-org/llama.cpp/issues/9291)
+[ggml](https://github.com/ggml-org/ggml) / [ops](https://github.com/ggml-org/llama.cpp/blob/master/docs/ops.md) / [maintainer PRs](https://github.com/ggml-org/llama.cpp/pulls?q=is%3Apr%20is%3Aopen%20draft%3AFalse%20(author%3Asanmai%20OR%20author%3Argerganov%20OR%20author%3AKitaitiMakoto%20OR%20author%3Adanbev%20OR%20author%3Aaldehir%20OR%20author%3Amax-krasnyansky%20OR%20author%3ACISC%20OR%20author%3Aggerganov%20OR%20author%3Aam17an%20OR%20author%3Ajhen0409%20OR%20author%3Abartowski1182%20OR%20author%3Anikwen%20OR%20author%3Ahipudding%20OR%20author%3Aravi9%20OR%20author%3AServeurpersoCom%20OR%20author%3Athelittlefireman%20OR%20author%3Apwilkin%20OR%20author%3Areeselevine%20OR%20author%3Angxson%20OR%20author%3Ajeffbolznv%20OR%20author%3Amarty1885%20OR%20author%3A0cc4m%20OR%20author%3ATitaniumtown%20OR%20author%3Aangt%20OR%20author%3AIMbackK%20OR%20author%3Aarthw%20OR%20author%3AJohannesGaessler%20OR%20author%3AORippler%20OR%20author%3Aruixiang63%20OR%20author%3Axctan%20OR%20author%3AmasterFoad%20OR%20author%3Aallozaur%20OR%20author%3AAndreas-Krebbel%20OR%20author%3Ayomaytk%20OR%20author%3Aaendk%20OR%20author%3Awine99%20OR%20author%3Agaugarg-nv%20OR%20author%3Ataronaeo%20OR%20author%3Aforforever73%20OR%20author%3Achaxu01%20OR%20author%3Alhez%20OR%20author%3AAlekseiNikiforovIBM%20OR%20author%3Anetrunnereve%20OR%20author%3Afairydreaming%20OR%20author%3Aynankani)%20sort%3Aupdated-desc) / [dev stats](https://github.com/ggml-org/llama.cpp-dev) / [lib llama API](https://github.com/ggml-org/llama.cpp/issues/9289) / [llama-server REST API](https://github.com/ggml-org/llama.cpp/issues/9291)
 
 </div>
 
@@ -79,7 +79,7 @@ The `llama.cpp` project is build on top of the [ggml](https://github.com/ggml-or
 | --- | --- |
 | [BLAS](docs/build.md#blas-build) | All |
 | [BLIS](docs/backend/BLIS.md) | All |
-| [CANN](docs/build.md#cann) | Ascend NPU |
+| [CANN](docs/build.md#cann) | Huawei Ascend NPU |
 | [CUDA](docs/build.md#cuda) | Nvidia GPU |
 | [HIP](docs/build.md#hip) | AMD GPU |
 | [Hexagon](docs/backend/snapdragon/README.md) | Snapdragon |

cmake/win-app.cmake

@@ -0,0 +1,66 @@
+set(LLAMA_WIN_APP_MANIFEST "${CMAKE_CURRENT_LIST_DIR}/../common/win-app.manifest")
+
+if(WIN32)
+    if(MSVC OR CMAKE_CXX_SIMULATE_ID STREQUAL "MSVC")
+        set(LLAMA_WIN_APP_MSVC_ABI TRUE)
+    else()
+        set(LLAMA_WIN_APP_MSVC_ABI FALSE)
+    endif()
+
+    if(NOT LLAMA_WIN_APP_MSVC_ABI)
+        if(CMAKE_VERSION VERSION_LESS 3.21)
+            message(FATAL_ERROR "CMake 3.21 or newer is required to build for Windows with the GNU ABI")
+        endif()
+        enable_language(RC)
+        set(LLAMA_WIN_APP_MANIFEST_RC "${CMAKE_CURRENT_BINARY_DIR}/win-app-manifest.rc")
+        configure_file("${CMAKE_CURRENT_LIST_DIR}/win-app.rc.in" "${LLAMA_WIN_APP_MANIFEST_RC}" @ONLY)
+        set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS "${LLAMA_WIN_APP_MANIFEST}")
+    endif()
+endif()
+
+function(llama_add_app_manifest target)
+    if(NOT WIN32)
+        return()
+    endif()
+    get_target_property(_aliased ${target} ALIASED_TARGET)
+    if(_aliased)
+        return()
+    endif()
+    get_target_property(_type ${target} TYPE)
+    if(NOT _type STREQUAL "EXECUTABLE")
+        return()
+    endif()
+    if(LLAMA_WIN_APP_MSVC_ABI)
+        target_sources(${target} PRIVATE "${LLAMA_WIN_APP_MANIFEST}")
+    else()
+        if(NOT TARGET llama-win-app-manifest-res)
+            add_library(llama-win-app-manifest-res OBJECT ${LLAMA_WIN_APP_MANIFEST_RC})
+            set_source_files_properties(${LLAMA_WIN_APP_MANIFEST_RC} PROPERTIES OBJECT_DEPENDS "${LLAMA_WIN_APP_MANIFEST}")
+        endif()
+        target_link_libraries(${target} PRIVATE "$<TARGET_OBJECTS:llama-win-app-manifest-res>")
+    endif()
+endfunction()
+
+function(llama_apply_app_manifests_dir dir)
+    get_property(_targets DIRECTORY "${dir}" PROPERTY BUILDSYSTEM_TARGETS)
+    foreach(_target IN LISTS _targets)
+        llama_add_app_manifest(${_target})
+    endforeach()
+    get_property(_subdirs DIRECTORY "${dir}" PROPERTY SUBDIRECTORIES)
+    foreach(_subdir IN LISTS _subdirs)
+        llama_apply_app_manifests_dir("${_subdir}")
+    endforeach()
+endfunction()
+
+function(llama_apply_app_manifests)
+    if(NOT WIN32)
+        return()
+    endif()
+    get_property(_subdirs DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}" PROPERTY SUBDIRECTORIES)
+    foreach(_dir IN LISTS _subdirs)
+        get_filename_component(_name "${_dir}" NAME)
+        if(_name STREQUAL "app" OR _name STREQUAL "examples" OR _name STREQUAL "pocs" OR _name STREQUAL "tests" OR _name STREQUAL "tools")
+            llama_apply_app_manifests_dir("${_dir}")
+        endif()
+    endforeach()
+endfunction()

cmake/win-app.rc.in

@@ -0,0 +1,3 @@
+// generated from cmake/win-app.rc.in - do not edit
+// RT_MANIFEST (24), ID 1: the application manifest shared by all llama.cpp executables
+1 24 "@LLAMA_WIN_APP_MANIFEST@"

common/arg.cpp

@@ -12,16 +12,6 @@
 #include "speculative.h"
 #include "preset.h"
 
-// fix problem with std::min and std::max
-#if defined(_WIN32)
-#define WIN32_LEAN_AND_MEAN
-#ifndef NOMINMAX
-#   define NOMINMAX
-#endif
-#include <windows.h>
-#include <shellapi.h>
-#endif
-
 #include <algorithm>
 #include <cinttypes>
 #include <climits>
@@ -38,20 +28,6 @@
 #include <thread> // for hardware_concurrency
 #include <vector>
 
-#ifndef __EMSCRIPTEN__
-#ifdef __linux__
-#include <linux/limits.h>
-#elif defined(_WIN32)
-#   if !defined(PATH_MAX)
-#   define PATH_MAX MAX_PATH
-#   endif
-#elif defined(_AIX)
-#include <sys/limits.h>
-#else
-#include <sys/syslimits.h>
-#endif
-#endif
-
 #define LLAMA_MAX_URL_LENGTH 2084 // Maximum URL Length in Chrome: 2083
 
 using json = common_json;
@@ -351,7 +327,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()).string();
 }
 
 static bool spec_types_is_default(const common_params & params) {
@@ -762,7 +738,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.string().c_str());
         common_preset global;
         common_presets presets = ctx.load_from_ini(path, global);
         global.apply_to_params(params);
@@ -1259,44 +1235,7 @@ bool common_params_to_map(int argc, char ** argv, llama_example ex, std::map<com
     return true;
 }
 
-#ifdef _WIN32
-struct utf8_argv {
-    std::vector<std::string> buf;
-    std::vector<char*> ptrs;
-};
-
-static utf8_argv make_utf8_argv() {
-    utf8_argv out;
-    int wargc = 0;
-    LPWSTR* wargv = CommandLineToArgvW(GetCommandLineW(), &wargc);
-    if (!wargv) return out;
-
-    out.buf.reserve(wargc);
-    for (int i = 0; i < wargc; ++i) {
-        int n = WideCharToMultiByte(CP_UTF8, WC_ERR_INVALID_CHARS, wargv[i], -1, nullptr, 0, nullptr, nullptr);
-        if (n <= 0) { out.buf.emplace_back(); continue; }
-        auto& s = out.buf.emplace_back();
-        s.resize(static_cast<size_t>(n - 1));
-        (void)WideCharToMultiByte(CP_UTF8, 0, wargv[i], -1, s.data(), n, nullptr, nullptr);
-    }
-    LocalFree(wargv);
-
-    out.ptrs.reserve(out.buf.size() + 1);
-    for (auto& s : out.buf) out.ptrs.push_back(s.data());
-    out.ptrs.push_back(nullptr);
-    return out;
-}
-#endif
-
 bool common_params_parse(int argc, char ** argv, common_params & params, llama_example ex, void(*print_usage)(int, char **)) {
-#ifdef _WIN32
-    auto utf8 = make_utf8_argv();
-    // repair argv only when it matches the process command line
-    if (static_cast<int>(utf8.buf.size()) == argc) {
-        argv = utf8.ptrs.data();
-    }
-#endif
-
     auto ctx_arg = common_params_parser_init(params, ex, print_usage);
     const common_params params_org = ctx_arg.params; // the example can modify the default params
 

common/common.cpp

@@ -219,6 +219,23 @@ int32_t common_cpu_get_num_math() {
             }
         }
     }
+#elif defined(__riscv) && defined(__linux__)
+    std::ifstream cpuinfo("/proc/cpuinfo");
+    std::string line;
+    int32_t n_a100 = 0;
+    while (std::getline(cpuinfo, line)) {
+        std::istringstream fields(line);
+        std::string key;
+        std::string separator;
+        uint64_t marchid;
+        if (fields >> key >> separator >> std::hex >> marchid &&
+            key == "marchid" && separator == ":" && marchid == UINT64_C(0x8000000041000002)) {
+            ++n_a100;
+        }
+    }
+    if (n_a100 > 0) {
+        return n_a100;
+    }
 #elif defined(__powerpc64__) || defined(__powerpc__)
     int32_t smt_factor = 1;
     int phy_cpus = common_cpu_get_num_physical_cores();
@@ -917,12 +934,6 @@ std::string wstring_to_utf8(const std::wstring & str) {
 }
 #endif
 
-// returns the path as a UTF-8 string, preserving its separators
-std::string fs_path_to_utf8(const std::filesystem::path & path) {
-    const auto value = path.u8string();
-    return std::string(value.begin(), value.end());
-}
-
 void fs_write_atomic(const std::filesystem::path & path, const std::string & data) {
     // a random suffix gives every writer its own temporary file, across processes and containers
     std::random_device rd;
@@ -946,7 +957,7 @@ void fs_write_atomic(const std::filesystem::path & path, const std::string & dat
 
     if (file.fail() || ec) {
         std::filesystem::remove(path_tmp, ec);
-        throw std::runtime_error("failed to write file: " + fs_path_to_utf8(path));
+        throw std::runtime_error("failed to write file: " + path.string());
     }
 }
 
@@ -1040,7 +1051,7 @@ std::filesystem::path fs_get_cache_file(const std::string & filename) {
     std::error_code ec;
     common_create_directories(cache_directory, ec);
     if (ec) {
-        throw std::runtime_error("failed to create cache directory: " + fs_path_to_utf8(cache_directory));
+        throw std::runtime_error("failed to create cache directory: " + cache_directory.string());
     }
     return cache_directory / std::filesystem::u8path(filename);
 }
@@ -2264,6 +2275,16 @@ int32_t common_batch::add_embd(llama_embd embd, const llama_pos * pos, llama_seq
     return size() - 1;
 }
 
+int32_t common_batch::add_embd_row(int32_t row, const llama_pos * pos, llama_seq_id seq_id, bool output) {
+    token t = { LLAMA_TOKEN_NULL, { 0, 0, 0, 0 }, seq_id, output, { nullptr, 0, 0 }, { nullptr, 0, 0 }, {} };
+    t.embd_row = row;
+    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());
@@ -2284,6 +2305,12 @@ llama_batch_ext * common_batch::get_sub_batch(int32_t off, int32_t n) {
             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 if (t.embd_row >= 0) {
+            idx = llama_batch_ext_add_embd_row(res, t.seq_id, t.embd_row);
+            if (idx < 0) {
+                GGML_ABORT("%s: failed to add embd row %d at index %d (error %d, n = %d)\n", __func__, t.embd_row, i, idx, n);
+            }
+            llama_batch_ext_set_pos(res, idx, t.pos.data());
         } else {
             idx = llama_batch_ext_add_embd(res, t.seq_id, t.embd);
             if (idx < 0) {

common/common.h

@@ -910,9 +910,6 @@ std::wstring utf8_to_wstring(const std::string & str);
 std::string  wstring_to_utf8(const std::wstring & str);
 #endif
 
-// returns the path as a UTF-8 string, preserving its separators
-std::string fs_path_to_utf8(const std::filesystem::path & path);
-
 //
 // Environment utils
 //
@@ -1099,6 +1096,7 @@ struct common_batch {
         llama_embd   state; // non-owning view of the data passed to set_embd_state(), data == NULL if none
         std::vector<llama_seq_id> seq_ids_extra; // see add_seq()
         int32_t      decision_order = 0; // see llama_batch_ext_set_decision_order()
+        int32_t      embd_row = -1; // see add_embd_row(), -1 if none
     };
 
     std::vector<token> tokens; // mirror of the entries, tokens[i] describes batch index i
@@ -1141,6 +1139,10 @@ struct common_batch {
     // pos points to n_pos positions
     int32_t add_embd(llama_embd embd, const llama_pos * pos, llama_seq_id seq_id, bool output);
 
+    // add an entry whose embedding is row `row` of the layer input device buffer of the context's ctx_other
+    // pos points to n_pos positions
+    int32_t add_embd_row(int32_t row, const llama_pos * pos, llama_seq_id seq_id, bool output);
+
     int32_t size() const { return (int32_t) tokens.size(); }
 };
 

common/download.cpp

@@ -1005,7 +1005,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).string();
 
         const std::string blob_url = url_prefix + "/blobs/" + gguf_digest;
         common_download_opts opts;

common/hf-cache.cpp

@@ -63,7 +63,7 @@ static fs::path get_cache_directory() {
 }
 
 std::string get_cache_path() {
-    return fs_path_to_utf8(get_cache_directory());
+    return get_cache_directory().string();
 }
 
 static std::string folder_name_to_repo(const std::string & folder) {
@@ -340,11 +340,11 @@ 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);
+            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;
             }
@@ -411,7 +411,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;
@@ -436,7 +436,7 @@ hf_files get_cached_files(const std::string & repo_id) {
                 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.local_path = entry.path().string();
                 file.final_path = file.local_path;
                 files.push_back(std::move(file));
             }
@@ -516,7 +516,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;

common/jinja/runtime.cpp

@@ -307,58 +307,13 @@ value binary_expression::execute_impl(context & ctx) const {
     throw std::runtime_error("Unknown operator \"" + op.value + "\" between " + left_val->type() + " and " + right_val->type());
 }
 
-static value try_builtin_func(context & ctx, const std::string & name, value & input, bool undef_on_missing = false) {
-    JJ_DEBUG("Trying built-in function '%s' for type %s", name.c_str(), input->type().c_str());
-    if (ctx.is_get_stats) {
-        value_t::stats_t::mark_used(input);
-        input->stats.ops.insert(name);
-    }
-    auto builtins = input->get_builtins();
-    auto it = builtins.find(name);
-    if (it != builtins.end()) {
-        JJ_DEBUG("Binding built-in '%s'", name.c_str());
-        return mk_val<value_func>(name, it->second, input);
-    }
-    if (undef_on_missing) {
-        return mk_val<value_undefined>(name);
-    }
-    throw std::runtime_error("Unknown (built-in) filter '" + name + "' for type " + input->type());
-}
-
 static value apply_filter(context & ctx, const statement_ptr & filter, value input) {
     JJ_DEBUG("Applying filter to %s", input->type().c_str());
 
-    auto set_filter_alias = [](auto & filter_id) {
-        if (filter_id == "count") {
-            filter_id = "length";
-        } else if (filter_id == "d") {
-            filter_id = "default";
-        } else if (filter_id == "e") {
-            filter_id = "escape";
-        } else if (filter_id == "trim") {
-            filter_id = "strip";
-        }
-    };
-
     if (is_stmt<identifier>(filter)) {
         auto filter_id = cast_stmt<identifier>(filter)->val;
-
-        set_filter_alias(filter_id);
         JJ_DEBUG("Applying filter '%s' to %s", filter_id.c_str(), input->type().c_str());
-        // TODO: Refactor filters so this coercion can be done automatically
-        if (!input->is_undefined() && !is_val<value_string>(input) && (
-            filter_id == "capitalize" ||
-            filter_id == "lower" ||
-            filter_id == "replace" ||
-            filter_id == "strip" ||
-            filter_id == "title" ||
-            filter_id == "upper" ||
-            filter_id == "wordcount"
-        )) {
-            JJ_DEBUG("Coercing %s to String for '%s' filter", input->type().c_str(), filter_id.c_str());
-            input = mk_val<value_string>(input->as_string());
-        }
-        return try_builtin_func(ctx, filter_id, input)->invoke(func_args(ctx));
+        return apply_filter_by_name(ctx, filter_id, input, func_args(ctx));
 
     } else if (is_stmt<call_expression>(filter)) {
         auto call = cast_stmt<call_expression>(filter);
@@ -367,14 +322,13 @@ static value apply_filter(context & ctx, const statement_ptr & filter, value inp
         }
         auto filter_id = cast_stmt<identifier>(call->callee)->val;
 
-        set_filter_alias(filter_id);
         JJ_DEBUG("Applying filter '%s' with arguments to %s", filter_id.c_str(), input->type().c_str());
         func_args args(ctx);
         for (const auto & arg_expr : call->args) {
             args.push_back(arg_expr->execute(ctx));
         }
 
-        return try_builtin_func(ctx, filter_id, input)->invoke(args);
+        return apply_filter_by_name(ctx, filter_id, input, args);
 
     } else {
         throw std::runtime_error("Invalid filter expression");

common/jinja/value.cpp

@@ -72,6 +72,17 @@ const std::vector<value> & func_args::get_args() const {
     return args;
 }
 
+void func_args::ensure_kwargs(std::initializer_list<const char *> names) const {
+    for (const auto & arg : args) {
+        if (is_val<value_kwarg>(arg)) {
+            auto * kwarg = cast_val<value_kwarg>(arg);
+            if (std::find(names.begin(), names.end(), kwarg->key) == names.end()) {
+                throw raised_exception("Function '" + func_name + "' got an unexpected keyword argument '" + kwarg->key + "'");
+            }
+        }
+    }
+}
+
 /**
  * Function that mimics Python's array slicing.
  */
@@ -134,6 +145,14 @@ static value empty_value_fn(const func_args &) {
         return mk_val<T>();
     }
 }
+
+template<typename T>
+static value empty_value_no_args_fn(const func_args & args) {
+    args.ensure_kwargs();
+    args.ensure_count(1, 1);
+    return empty_value_fn<T>(args);
+}
+
 template<typename T>
 static value test_type_fn(const func_args & args) {
     args.ensure_count(1);
@@ -594,6 +613,52 @@ const func_builtins & global_builtins() {
     return builtins;
 }
 
+value try_builtin_func(context & ctx, const std::string & name, value & input, bool undef_on_missing) {
+    JJ_DEBUG("Trying built-in function '%s' for type %s", name.c_str(), input->type().c_str());
+    if (ctx.is_get_stats) {
+        value_t::stats_t::mark_used(input);
+        input->stats.ops.insert(name);
+    }
+    auto builtins = input->get_builtins();
+    auto it = builtins.find(name);
+    if (it != builtins.end()) {
+        JJ_DEBUG("Binding built-in '%s'", name.c_str());
+        return mk_val<value_func>(name, it->second, input);
+    }
+    if (undef_on_missing) {
+        return mk_val<value_undefined>(name);
+    }
+    throw std::runtime_error("Unknown (built-in) filter '" + name + "' for type " + input->type());
+}
+
+value apply_filter_by_name(context & ctx, std::string filter_id, value input, const func_args & args) {
+    if (filter_id == "count") {
+        filter_id = "length";
+    } else if (filter_id == "d") {
+        filter_id = "default";
+    } else if (filter_id == "e") {
+        filter_id = "escape";
+    } else if (filter_id == "trim") {
+        filter_id = "strip";
+    }
+
+    // TODO: Refactor filters so this coercion can be done automatically
+    if (!input->is_undefined() && !is_val<value_string>(input) && (
+        filter_id == "capitalize" ||
+        filter_id == "lower" ||
+        filter_id == "replace" ||
+        filter_id == "strip" ||
+        filter_id == "title" ||
+        filter_id == "upper" ||
+        filter_id == "wordcount"
+    )) {
+        JJ_DEBUG("Coercing %s to String for '%s' filter", input->type().c_str(), filter_id.c_str());
+        input = mk_val<value_string>(input->as_string());
+    }
+
+    return try_builtin_func(ctx, filter_id, input)->invoke(args);
+}
+
 
 const func_builtins & value_int_t::get_builtins() const {
     static const func_builtins builtins = {
@@ -667,16 +732,22 @@ const func_builtins & value_string_t::get_builtins() const {
     static const func_builtins builtins = {
         {"default", default_value},
         {"upper", [](const func_args & args) -> value {
+            args.ensure_kwargs();
+            args.ensure_count(1, 1);
             args.ensure_vals<value_string>();
             jinja::string str = args.get_pos(0)->as_string().uppercase();
             return mk_val<value_string>(str);
         }},
         {"lower", [](const func_args & args) -> value {
+            args.ensure_kwargs();
+            args.ensure_count(1, 1);
             args.ensure_vals<value_string>();
             jinja::string str = args.get_pos(0)->as_string().lowercase();
             return mk_val<value_string>(str);
         }},
         {"strip", [](const func_args & args) -> value {
+            args.ensure_kwargs({"chars"});
+            args.ensure_count(1, 2);
             value val_input = args.get_pos(0);
             if (!is_val<value_string>(val_input)) {
                 throw raised_exception("strip() first argument must be a string");
@@ -707,11 +778,15 @@ const func_builtins & value_string_t::get_builtins() const {
             }
         }},
         {"title", [](const func_args & args) -> value {
+            args.ensure_kwargs();
+            args.ensure_count(1, 1);
             args.ensure_vals<value_string>();
             jinja::string str = args.get_pos(0)->as_string().titlecase();
             return mk_val<value_string>(str);
         }},
         {"capitalize", [](const func_args & args) -> value {
+            args.ensure_kwargs();
+            args.ensure_count(1, 1);
             args.ensure_vals<value_string>();
             jinja::string str = args.get_pos(0)->as_string().capitalize();
             return mk_val<value_string>(str);
@@ -1013,6 +1088,10 @@ const func_builtins & value_bool_t::get_builtins() const {
     return builtins;
 }
 
+[[noreturn]] static value array_sum_not_implemented(const func_args &) {
+    throw not_implemented_exception("Array sum builtin not implemented");
+}
+
 const func_builtins & value_array_t::get_builtins() const {
     static const func_builtins builtins = {
         {"default", default_value},
@@ -1079,7 +1158,9 @@ const func_builtins & value_array_t::get_builtins() const {
         {"select", select_or_reject<false, false>},
         {"rejectattr", select_or_reject<true, true>},
         {"reject", select_or_reject<true, false>},
+        {"sum", array_sum_not_implemented},
         {"join", [](const func_args & args) -> value {
+            args.ensure_kwargs({"d", "attribute"});
             args.ensure_count(1, 3);
             if (!is_val<value_array>(args.get_pos(0))) {
                 throw raised_exception("join() first argument must be an array");
@@ -1117,13 +1198,28 @@ const func_builtins & value_array_t::get_builtins() const {
         {"tojson", tojson},
         {"map", [](const func_args & args) -> value {
             args.ensure_count(2);
-            if (!is_val<value_array>(args.get_pos(0))) {
+            value val = args.get_pos(0);
+            if (!is_val<value_array>(val)) {
                 throw raised_exception("map: first argument must be an array");
             }
             if (!is_val<value_kwarg>(args.get_args().at(1))) {
-                throw not_implemented_exception("map: filter-mapping not implemented");
+                value filter = args.get_pos(1);
+                if (!is_val<value_string>(filter)) {
+                    throw raised_exception("map: filter name must be a string");
+                }
+                auto out = mk_val<value_array>();
+                func_args filter_args(args.ctx);
+                const auto & all_args = args.get_args();
+                for (size_t i = 2; i < all_args.size(); ++i) {
+                    filter_args.push_back(all_args[i]);
+                }
+                const std::string filter_id = filter->as_string().str();
+                auto arr = val->as_array();
+                for (const auto & item : arr) {
+                    out->push_back(apply_filter_by_name(args.ctx, filter_id, item, filter_args));
+                }
+                return is_val<value_tuple>(val) ? mk_val<value_tuple>(out->as_array()) : out;
             }
-            value val       = args.get_pos(0);
             value attribute = args.get_kwarg_or_pos("attribute", 1);
             value default_val = args.get_kwarg("default", mk_val<value_undefined>());
             auto out = mk_val<value_array>();
@@ -1132,7 +1228,7 @@ const func_builtins & value_array_t::get_builtins() const {
                 value attr_val = get_attribute(item, attribute, default_val);
                 out->push_back(attr_val);
             }
-            return is_val<value_tuple>(val) ? mk_val<value_tuple>(std::move(out->as_array())) : out;
+            return is_val<value_tuple>(val) ? mk_val<value_tuple>(out->as_array()) : out;
         }},
         {"append", [](const func_args & args) -> value {
             args.ensure_count(2);
@@ -1405,14 +1501,18 @@ const func_builtins & value_none_t::get_builtins() const {
 const func_builtins & value_undefined_t::get_builtins() const {
     static const func_builtins builtins = {
         {"default", default_value},
-        {"capitalize", empty_value_fn<value_string>},
+        {"capitalize", empty_value_no_args_fn<value_string>},
         {"first", empty_value_fn<value_undefined>},
         {"items", empty_value_fn<value_array>},
-        {"join", empty_value_fn<value_string>},
+        {"join", [](const func_args & args) -> value {
+            args.ensure_kwargs({"d", "attribute"});
+            args.ensure_count(1, 3);
+            return empty_value_fn<value_string>(args);
+        }},
         {"last", empty_value_fn<value_undefined>},
         {"length", empty_value_fn<value_int>},
         {"list", empty_value_fn<value_array>},
-        {"lower", empty_value_fn<value_string>},
+        {"lower", empty_value_no_args_fn<value_string>},
         {"map", empty_value_fn<value_array>},
         {"max", empty_value_fn<value_undefined>},
         {"min", empty_value_fn<value_undefined>},
@@ -1425,12 +1525,16 @@ const func_builtins & value_undefined_t::get_builtins() const {
         {"selectattr", empty_value_fn<value_array>},
         {"sort", empty_value_fn<value_array>},
         {"string", empty_value_fn<value_string>},
-        {"strip", empty_value_fn<value_string>},
+        {"strip", [](const func_args & args) -> value {
+            args.ensure_kwargs({"chars"});
+            args.ensure_count(1, 2);
+            return empty_value_fn<value_string>(args);
+        }},
         {"sum", empty_value_fn<value_int>},
-        {"title", empty_value_fn<value_string>},
+        {"title", empty_value_no_args_fn<value_string>},
         {"truncate", empty_value_fn<value_string>},
         {"unique", empty_value_fn<value_array>},
-        {"upper", empty_value_fn<value_string>},
+        {"upper", empty_value_no_args_fn<value_string>},
         {"wordcount", empty_value_fn<value_int>},
     };
     return builtins;

common/jinja/value.h

@@ -7,6 +7,7 @@
 #include <cmath>
 #include <cstdint>
 #include <functional>
+#include <initializer_list>
 #include <map>
 #include <memory>
 #include <set>
@@ -176,6 +177,8 @@ protected:
 //
 
 const func_builtins & global_builtins();
+value try_builtin_func(context & ctx, const std::string & name, value & input, bool undef_on_missing = false);
+value apply_filter_by_name(context & ctx, std::string filter_id, value input, const func_args & args);
 
 std::string value_to_json(const value & val, int indent = -1, const std::string_view item_sep = ", ", const std::string_view key_sep = ": ");
 
@@ -658,6 +661,7 @@ public:
     size_t count() const { return args.size(); }
     void push_back(const value & val);
     void push_front(const value & val);
+    void ensure_kwargs(std::initializer_list<const char *> names = {}) const;
     void ensure_count(size_t min, size_t max = 999) const {
         size_t n = args.size();
         if (n < min || n > max) {

common/parsers/gpt-oss.cpp

@@ -71,6 +71,7 @@ common_chat_params common_chat_params_init_gpt_oss(const common_chat_template &
         auto content         = p.rule("message-content", p.until("<|end|>"));
         auto channel         = p.literal("<|channel|>") + (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|>")) + constrain_type);
 
         // Occasionally, gpt-oss-20b will prefix channels with this commentary
         auto stray_commentary = p.optional(p.literal("<|channel|>commentary") + p.optional(p.literal(" to=assistant")));
@@ -84,7 +85,7 @@ common_chat_params common_chat_params_init_gpt_oss(const common_chat_template &
 
         auto analysis = p.ref("analysis");
         auto preamble = p.rule("preamble", p.literal("<|channel|>commentary<|message|>") + p.content(content) + end);
-        auto final_msg = p.rule("final", stray_commentary + p.literal("<|channel|>final<|message|>") + p.content(content));
+        auto final_msg = p.rule("final", stray_commentary + p.literal("<|channel|>final") + constraint + p.literal("<|message|>") + p.content(content));
 
         // Consume any unsolicited tool calls, e.g. builtin functions
         auto unsolicited = p.rule("unsolicited", p.atomic(p.optional(channel) + p.literal(" to=") + content + end));
@@ -92,7 +93,6 @@ common_chat_params common_chat_params_init_gpt_oss(const common_chat_template &
         auto any = p.rule("any", preamble | analysis);
 
         if (has_response_format) {
-            auto constraint = p.optional(p.space() + p.optional(p.literal("<|constrain|>")) + constrain_type);
             auto response_format = p.rule("response-format",
                 p.literal("<|channel|>final") + constraint + p.literal("<|message|>") +
                 p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema)));
@@ -109,7 +109,6 @@ common_chat_params common_chat_params_init_gpt_oss(const common_chat_template &
                 const auto   params   = common_chat_tool_parameters(function);
 
                 auto func_name  = p.literal(" to=functions.") + p.tool_name(p.literal(name));
-                auto constraint = p.optional(p.space() + p.optional(p.literal("<|constrain|>")) + constrain_type);
                 auto args       = p.tool_args(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-schema", params));
 
                 // recipient in role header

common/parsers/qwen3-coder.cpp

@@ -74,16 +74,24 @@ common_chat_params common_chat_params_init_qwen3_coder(const common_chat_templat
     data.parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
         auto generation_prompt = p.literal(GEN_PREFIX);
 
+        // The template strips the newline before </think>, so leave it out of the reasoning
+        auto think_body  = p.until_one_of({ "\n</think>", "</think>", "<tool_call>" });
+        auto think_close = p.literal("\n</think>") | p.literal("</think>") | p.peek(p.literal("<tool_call>"));
+
+        auto reasoning_extracted = p.optional("<think>" + p.space() + p.reasoning(think_body) + think_close);
+        auto reasoning_skipped   = p.optional("<think>" + p.space() + think_body + think_close);
+
         auto reasoning = p.eps();
         if (supports_reasoning && extract_reasoning) {
-            // The template strips the newline before </think>, so leave it out of the reasoning
-            reasoning = p.optional("<think>" + p.space() +
-                                   p.reasoning(p.until_one_of({ "\n</think>", "</think>", "<tool_call>" })) +
-                                   (p.literal("\n</think>") | p.literal("</think>") | p.peek(p.literal("<tool_call>"))));
+            reasoning = reasoning_extracted;
         }
 
         // Response format parser
         if (has_response_format) {
+            // The answer must be bare JSON, so the think block is consumed even when it is not extracted
+            if (supports_reasoning && !extract_reasoning) {
+                reasoning = reasoning_skipped;
+            }
             return generation_prompt + (reasoning << p.content(p.schema(p.json(), "response-format", inputs.json_schema)));
         }
 

common/preset.cpp

@@ -171,12 +171,12 @@ static std::map<std::string, std::map<std::string, std::string>> parse_ini_from_
     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.string());
     }
 
     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.string());
     }
 
     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.string());
     }
 
     std::string current_section = COMMON_PRESET_DEFAULT_NAME;
@@ -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.string().c_str());
             } else {
                 throw std::runtime_error(string_format(
                     "option '%s' not recognized in preset '%s'",
@@ -401,7 +401,7 @@ common_presets common_preset_context::load_from_models_dir(const std::string & m
             if (!entry.is_regular_file(ec)) {
                 continue;
             }
-            const std::string fname = fs_path_to_utf8(entry.path().filename());
+            const std::string fname = entry.path().filename().string();
             if (!string_ends_with(fname, ".gguf")) {
                 continue;
             }
@@ -422,9 +422,9 @@ common_presets common_preset_context::load_from_models_dir(const std::string & m
         if (!path.empty()) {
             models.push_back({
                 /* name        */ name,
-                /* path        */ fs_path_to_utf8(path),
-                /* path_mmproj */ fs_path_to_utf8(mmproj_file), // can be empty
-                /* path_draft  */ fs_path_to_utf8(draft_file)    // can be empty
+                /* path        */ path.string(),
+                /* path_mmproj */ mmproj_file.string(), // can be empty
+                /* path_draft  */ draft_file.string()   // can be empty
             });
         }
     };
@@ -432,13 +432,13 @@ common_presets common_preset_context::load_from_models_dir(const std::string & m
     for (const auto & entry : std::filesystem::directory_iterator(dir)) {
         std::error_code ec;
         if (entry.is_directory(ec)) {
-            scan_subdir(entry.path(), fs_path_to_utf8(entry.path().filename()));
+            scan_subdir(entry.path(), entry.path().filename().string());
             continue;
         }
         if (!entry.is_regular_file(ec)) {
             continue;
         }
-        const std::string fname = fs_path_to_utf8(entry.path().filename());
+        const std::string fname = entry.path().filename().string();
         if (!string_ends_with(fname, ".gguf")) {
             continue;
         }
@@ -450,7 +450,7 @@ common_presets common_preset_context::load_from_models_dir(const std::string & m
         string_replace_all(name, ".gguf", "");
         models.push_back({
             /* name        */ name,
-            /* path        */ fs_path_to_utf8(entry.path()),
+            /* path        */ entry.path().string(),
             /* path_mmproj */ "",
             /* path_draft  */ ""
         });

common/reasoning-budget.cpp

@@ -136,7 +136,8 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
             // track the end sequence within forced_tokens so it is also reported on DONE
             const int32_t match = ctx->end_matcher.advance(token);
             ctx->force_pos++;
-            if (ctx->force_pos >= ctx->forced_tokens.size()) {
+            // an end tag that is already in the prompt (e.g. "<think></think>") also ends forcing
+            if (match >= 0 || ctx->force_pos >= ctx->forced_tokens.size()) {
                 ctx->state = REASONING_BUDGET_DONE;
                 ctx->end_match = match;
                 COM_TRC("%s", "forced sequence complete, done\n");

common/speculative.cpp

@@ -1027,6 +1027,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
     const int32_t * target_layer_ids   = nullptr; // model_dft's extract layer indices
     uint32_t        target_layer_ids_n = 0;
 
+    // inject target features from the target's device buffer instead of host copies (env LLAMA_DFLASH_DEV=0 disables)
+    bool use_dev_features = false;
+
     common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq,
             common_speculative_type type = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH)
         : common_speculative_impl(type, n_seq, params.draft.n_max)
@@ -1141,6 +1144,23 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
             llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true);
         }
 
+        // the device buffer keeps layers in ascending id order, the fc input expects target_layer_ids order
+        {
+            const char * env = getenv("LLAMA_DFLASH_DEV");
+            const bool ascending = std::adjacent_find(target_layer_ids, target_layer_ids + target_layer_ids_n,
+                    [](int32_t a, int32_t b) { return a >= b; }) == target_layer_ids + target_layer_ids_n;
+
+            use_dev_features = !is_lilicorr && ascending && (env == nullptr || atoi(env) != 0);
+            if (use_dev_features) {
+                llama_set_embeddings_layer_inp_dev(ctx_tgt, true);
+                use_dev_features = llama_supports_embd_row(ctx_dft);
+                if (!use_dev_features) {
+                    llama_set_embeddings_layer_inp_dev(ctx_tgt, false);
+                }
+            }
+            LOG_INF("%s: DFlash target features via %s\n", __func__, use_dev_features ? "device buffer" : "host copies");
+        }
+
         // DFlash2 reads its selector lattice from h_nextn and never consumes raw logits.
         llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ !is_dflash2 && !is_lilicorr);
         llama_set_causal_attn(ctx_dft, causal_attn); // DFlash needs non-causal attention unless the model says otherwise
@@ -1233,6 +1253,25 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
             for (int32_t offset = 0; offset < n_rows; offset += n_ubatch) {
                 const int32_t n_chunk = std::min(n_ubatch, n_rows - offset);
 
+                if (use_dev_features) {
+                    // the features are the target batch rows in the target's device buffer
+                    batch_inject.clear();
+                    for (int32_t i = 0; i < n_chunk; ++i) {
+                        const int32_t   row        = i_batch_beg[seq_id] + offset + i;
+                        const llama_pos p          = batch_in.tokens[row].pos[0];
+                        const llama_pos pos_arr[4] = { p, p, p, 0 };
+                        batch_inject.add_embd_row(row, pos_arr, seq_id, false);
+                    }
+
+                    const int32_t rc = llama_process(ctx_dft, LLAMA_PROCESS_TYPE_DECODE, batch_inject.get());
+                    if (rc != 0) {
+                        LOG_ERR("%s: llama_process(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
+                                __func__, rc, (int) n_chunk, (int) offset);
+                        return false;
+                    }
+                    continue;
+                }
+
                 // 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);

common/win-app.manifest

@@ -0,0 +1,9 @@
+<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
+<assembly xmlns="urn:schemas-microsoft-com:asm.v1" manifestVersion="1.0">
+  <application xmlns="urn:schemas-microsoft-com:asm.v3">
+    <windowsSettings>
+      <activeCodePage xmlns="http://schemas.microsoft.com/SMI/2019/WindowsSettings">UTF-8</activeCodePage>
+      <longPathAware xmlns="http://schemas.microsoft.com/SMI/2016/WindowsSettings">true</longPathAware>
+    </windowsSettings>
+  </application>
+</assembly>

conversion/__init__.py

@@ -46,6 +46,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "ChatGLMForConditionalGeneration": "chatglm",
     "ChatGLMModel": "chatglm",
     "ClefModel": "clef",
+    "CloverMoeModelLM": "clover",
     "CodeShellForCausalLM": "codeshell",
     "CogVLMForCausalLM": "cogvlm",
     "Cohere2MoeForCausalLM": "command_r",
@@ -176,6 +177,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "Lfm3ForCausalLM": "lfm3",
     "Lfm3OmniForConditionalGeneration": "lfm3",
     "Lfm3VlForConditionalGeneration": "lfm3",
+    "LimiteForCausalLM": "limite",
     "Llama4ForCausalLM": "llama",
     "Llama4ForConditionalGeneration": "llama",
     "LlamaBidirectionalModel": "llama",
@@ -212,6 +214,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "Ministral3Model": "mistral3",
     "Mistral3ForConditionalGeneration": "mistral3",
     "Mistral3Model": "mistral3",
+    "MistralLarge4ForConditionalGeneration": "mistral4",
     "VoxtralRealtimeForConditionalGeneration": "mistral3",
     "MistralForCausalLM": "llama",
     "MixtralForCausalLM": "llama",
@@ -220,6 +223,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "ModernBertForSequenceClassification": "bert",
     "ModernBertModel": "bert",
     "NanbeigeForCausalLM": "nanbeige",
+    "Nemotron3_5TranscribeForConditionalGeneration": "nemotron",
     "NemotronForCausalLM": "nemotron",
     "NemotronHForCausalLM": "nemotron",
     "NemotronHPuzzleForCausalLM": "nemotron",
@@ -377,6 +381,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
     "MiniCPMOcrForConditionalGeneration": "minicpm",
     "Mistral3ForConditionalGeneration": "llava",
     "Mistral3Model": "llava",
+    "Nemotron3_5TranscribeForConditionalGeneration": "nemotron",
     "NemotronH_Omni_Reasoning_V3": "nemotron",
     "NemotronH_Nano_VL_V2": "nemotron",
     "MuseGlimmerForConditionalGeneration": "muse_glimmer",

conversion/base.py

@@ -16,6 +16,7 @@ from pathlib import Path
 from hashlib import sha256
 from typing import TYPE_CHECKING, Any, Callable, ContextManager, Iterable, Iterator, Literal, Sequence, TypeVar, cast
 from itertools import chain
+from math import prod
 from transformers import AutoConfig
 
 import numpy as np
@@ -155,6 +156,7 @@ class ModelBase:
         self.fuse_gate_up_exps = fuse_gate_up_exps
         self._gate_exp_buffer: dict[int, Tensor] = {}
         self._up_exp_buffer: dict[int, Tensor] = {}
+        self._unfusable_gate_up_layers: set[int] = set()
         self.fuse_qkv = fuse_qkv
         self._q_buffer: dict[int, Tensor] = {}
         self._k_buffer: dict[int, Tensor] = {}
@@ -175,7 +177,9 @@ class ModelBase:
         self._nvfp4_layer_algo: dict[str, str | None] = {} # per-layer quant_algo, keyed by HF module path
         self._prec_a4: dict[str, bool] = {} # gguf tensor name -> can use 4-bit (A4) activations
         self._fp8_as_q8 = fp8_as_q8
+        self._preserve_fp8 = ftype == gguf.LlamaFileType.GUESSED and not fp8_as_q8
         self._fp8_dequantized: set[str] = set()
+        self._fp8_e4m3_preserved: set[str] = set()
 
         # Apply heuristics to figure out typical tensor encoding based on first tensor's dtype
         # NOTE: can't use field "torch_dtype" in config.json, because some finetunes lie.
@@ -326,7 +330,7 @@ class ModelBase:
 
     def dequant_model(self):
         # If all quantized tensors were already handled (e.g. pure NVFP4), skip
-        if self._is_nvfp4 and not any(k.endswith((".weight_scale", ".weight_scale_inv")) for k in self.model_tensors):
+        if self._is_nvfp4 and not any(k.endswith((".weight_scale", ".weight_scale_inv", ".input_scale", ".activation_scale", "_activation_scale", ".k_scale", ".v_scale")) for k in self.model_tensors):
             return
 
         tensors_to_remove: list[str] = []
@@ -350,6 +354,12 @@ class ModelBase:
             def dequant_simple(weight: Tensor, scale: Tensor, block_size: Sequence[int] | None = None) -> Tensor:
                 scale = scale.float()
 
+                # ModelOpt block scales are pre-shaped for broadcasting, e.g. [out/128, 1, in/128, 1]
+                if block_size is None and scale.ndim == 2 * weight.ndim and all(n == 1 for n in scale.shape[1::2]):
+                    assert all(w % n == 0 for w, n in zip(weight.shape, scale.shape[0::2]))
+                    block_size = [w // n for w, n in zip(weight.shape, scale.shape[0::2])]
+                    scale = scale.reshape(tuple(scale.shape[0::2]))
+
                 if block_size is not None:
                     dim_offset = scale.ndim - len(block_size)
                     for i, size in enumerate(block_size):
@@ -440,7 +450,7 @@ class ModelBase:
 
                 return (unpacked * scale.unsqueeze(-1).float()).reshape(shape)
 
-            def dequant_fp8() -> None:
+            def dequant_fp8(fp8_configs: list[dict[str, Any]] | None = None) -> None:
                 for name in self.model_tensors.keys():
                     if name.endswith(".weight_scale"):
                         weight_name = name.removesuffix("_scale")
@@ -449,14 +459,29 @@ class ModelBase:
                             continue
                         w = self.model_tensors[weight_name]
                         s = self.model_tensors[name]
+                        block_size = None
                         is_fp8_weight = False
-                        if self._fp8_as_q8:
-                            is_fp8_weight = w().dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
-                        self.model_tensors[weight_name] = lambda w=w, s=s: dequant_simple(w(), s(), None)
+                        if self._fp8_as_q8 or fp8_configs is not None:
+                            weight = w()
+                            is_fp8_weight = weight.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
+                        if fp8_configs is not None:
+                            if not is_fp8_weight:
+                                raise ValueError(f"Expected remaining weight {weight_name!r} to be FP8, got {weight.dtype}")
+                            scale = s()
+                            is_channel_scale = scale.ndim >= 1 and scale.ndim <= weight.ndim and scale.shape[0] == weight.shape[0] and all(size == 1 for size in scale.shape[1:])
+                            if prod(scale.shape) != 1 and not is_channel_scale:
+                                block_configs = [c for c in fp8_configs if c.get("strategy") == "block"]
+                                if not block_configs or any(not c.get("block_structure") or c.get("group_size") is not None for c in block_configs):
+                                    raise ValueError(f"FP8 weight {weight_name!r} with scale shape {list(scale.shape)} requires a block layout")
+                                block_sizes = {tuple(c["block_structure"]) for c in block_configs}
+                                if len(block_sizes) != 1:
+                                    raise ValueError(f"FP8 weight {weight_name!r} has ambiguous block layouts")
+                                block_size = next(iter(block_sizes))
+                        self.model_tensors[weight_name] = lambda w=w, s=s, bs=block_size: dequant_simple(w(), s(), bs)
                         tensors_to_remove.append(name)
-                        if is_fp8_weight:
+                        if self._fp8_as_q8 and is_fp8_weight:
                             self._fp8_dequantized.add(weight_name)
-                    if name.endswith((".input_scale", ".k_scale", ".v_scale")):
+                    if name.endswith((".input_scale", ".activation_scale", "_activation_scale", ".k_scale", ".v_scale")):
                         tensors_to_remove.append(name)
 
             if quant_method == "bitnet":
@@ -521,29 +546,37 @@ class ModelBase:
                 nvfp4_compressed_tensors = self._is_nvfp4_compressed_tensors(quant_method, quant_format, groups)
 
                 if nvfp4_compressed_tensors:
-                    dequant_fp8()
+                    fp8_configs = [
+                        g["weights"] for g in groups.values()
+                        if isinstance(g, dict) and g.get("format") == "float-quantized"
+                        and g.get("weights", {}).get("type") == "float"
+                        and g.get("weights", {}).get("num_bits") == 8
+                    ]
+                    dequant_fp8(fp8_configs)
                 elif len(groups) > 1:
                     raise NotImplementedError("Can't handle multiple config groups for compressed-tensors yet")
                 elif quant_format == "float-quantized" or quant_format == "int-quantized" or quant_format == "naive-quantized":
                     weight_config = tuple(groups.values())[0]["weights"]
                     block_size = weight_config.get("block_structure", None)
                     strategy = weight_config.get("strategy")
-                    assert strategy == "channel" or strategy == "block"
-                    assert weight_config.get("group_size") is None  # didn't find a model using this yet
                     is_fp8 = (
                         quant_format == "float-quantized"
                         and weight_config.get("type") == "float"
                         and weight_config.get("num_bits") == 8
                     )
+                    assert strategy in ("channel", "block") or (is_fp8 and strategy == "tensor")
+                    assert weight_config.get("group_size") is None  # didn't find a model using this yet
                     for name in self.model_tensors.keys():
                         if name.endswith(".weight_scale"):
                             weight_name = name.removesuffix("_scale")
                             w = self.model_tensors[weight_name]
                             s = self.model_tensors[name]
-                            self.model_tensors[weight_name] = lambda w=w, s=s: dequant_simple(w(), s(), block_size)
+                            self.model_tensors[weight_name] = lambda w=w, s=s, bs=block_size: dequant_simple(w(), s(), bs)
                             tensors_to_remove.append(name)
                             if self._fp8_as_q8 and is_fp8:
                                 self._fp8_dequantized.add(weight_name)
+                        if is_fp8 and name.endswith((".input_scale", ".activation_scale", "_activation_scale", ".k_scale", ".v_scale")):
+                            tensors_to_remove.append(name)
                 elif quant_format == "pack-quantized":
                     weight_config = tuple(groups.values())[0]["weights"]
                     assert weight_config.get("strategy") == "group"
@@ -581,6 +614,148 @@ class ModelBase:
         for name, value in new_tensors.items():
             self.model_tensors[name] = value
 
+    def _map_fp8_weight_names(self, name: str) -> tuple[str, ...]:
+        new_name = self.tensor_map.get_name(name, try_suffixes=(".weight",))
+        return (new_name,) if new_name is not None else ()
+
+    def _prepare_fp8_e4m3_tensors(self):
+        if not self._preserve_fp8:
+            return
+
+        scale_tensors: dict[str, list[tuple[int, float]] | np.ndarray] = {}
+        input_scale_tensors: dict[str, list[tuple[int, float]] | np.ndarray] = {}
+        consumed: list[str] = []
+        candidates: dict[str, tuple[str, Tensor, tuple[str, ...], torch.Size]] = {}
+        weight_dtypes: dict[str, torch.dtype] = {}
+        scaled_weights: set[str] = set()
+        expert_pattern = re.compile(r"(\.[^.]*experts)\.(\d+)\.")
+        scale_types = {
+            gguf.MODEL_TENSOR.ATTN_Q, gguf.MODEL_TENSOR.ATTN_K, gguf.MODEL_TENSOR.ATTN_V,
+            gguf.MODEL_TENSOR.ATTN_OUT, gguf.MODEL_TENSOR.ATTN_QKV, gguf.MODEL_TENSOR.ATTN_GATE,
+            gguf.MODEL_TENSOR.FFN_GATE, gguf.MODEL_TENSOR.FFN_UP, gguf.MODEL_TENSOR.FFN_DOWN,
+            gguf.MODEL_TENSOR.FFN_GATE_SHEXP, gguf.MODEL_TENSOR.FFN_UP_SHEXP, gguf.MODEL_TENSOR.FFN_DOWN_SHEXP,
+            gguf.MODEL_TENSOR.FFN_GATE_EXP, gguf.MODEL_TENSOR.FFN_UP_EXP, gguf.MODEL_TENSOR.FFN_DOWN_EXP,
+            gguf.MODEL_TENSOR.SSM_IN, gguf.MODEL_TENSOR.SSM_OUT, gguf.MODEL_TENSOR.SSM_ALPHA, gguf.MODEL_TENSOR.SSM_BETA,
+            gguf.MODEL_TENSOR.NEXTN_EH_PROJ, gguf.MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, gguf.MODEL_TENSOR.OUTPUT,
+        }
+
+        for scale_name in list(self.model_tensors.keys()):
+            weight_name = None
+            if scale_name.endswith(".weight_scale"):
+                weight_name = scale_name.removesuffix("_scale")
+            elif scale_name.endswith("_scale_inv"):
+                weight_name = scale_name.removesuffix("_scale_inv")
+            elif scale_name.endswith(".qscale_weight"):
+                weight_name = scale_name.removesuffix("qscale_weight") + "weight"
+
+            if weight_name is None or weight_name not in self.model_tensors:
+                continue
+
+            scaled_weights.add(weight_name)
+            weight = self.model_tensors[weight_name]()
+            weight_dtypes[weight_name] = weight.dtype
+            if weight.dtype != torch.float8_e4m3fn or weight.ndim < 2:
+                continue
+
+            scale = self.model_tensors[scale_name]()
+            if prod(scale.shape) != 1:
+                continue
+
+            merged_name = expert_pattern.sub(r"\1.", weight_name)
+            new_names = self._map_fp8_weight_names(merged_name)
+            if not new_names:
+                continue
+            if any((mapped := self.tensor_map.get_type_and_name(new_name, try_suffixes=(".weight",))) is None or mapped[0] not in scale_types for new_name in new_names):
+                continue
+            bid = next((int(part) for part in new_names[0].split(".") if part.isdecimal()), None)
+            self._fp8_e4m3_preserved.add(weight_name)
+            if any(new_name.endswith("_norm.weight") or self.tensor_force_quant(weight_name, new_name, bid, weight.ndim) != gguf.GGMLQuantizationType.F8_E4M3 for new_name in new_names):
+                self._fp8_e4m3_preserved.discard(weight_name)
+                continue
+            candidates[weight_name] = (scale_name, scale, new_names, weight.shape)
+
+        expert_groups: dict[str, list[str]] = {}
+        for name in self.model_tensors:
+            if (match := expert_pattern.search(name)) and name.endswith(".weight"):
+                expert_groups.setdefault(name[:match.end(1)], []).append(name)
+        for names in expert_groups.values():
+            for name in names:
+                if name not in weight_dtypes:
+                    weight_dtypes[name] = self.model_tensors[name]().dtype
+            dtypes = {weight_dtypes[name] for name in names}
+            if torch.float8_e4m3fn not in dtypes:
+                continue
+            if dtypes == {torch.float8_e4m3fn} and not any(name in scaled_weights for name in names):
+                continue
+            # The converter merges expert projections together.
+            if not all(name in candidates for name in names):
+                for name in names:
+                    candidates.pop(name, None)
+                    self._fp8_e4m3_preserved.discard(name)
+                    if name not in scaled_weights:
+                        gen = self.model_tensors[name]
+                        self.model_tensors[name] = lambda gen=gen: gen().float()
+
+        for weight_name, (scale_name, scale, new_names, shape) in candidates.items():
+            scale = LazyTorchTensor.to_eager(scale).float().flatten()
+
+            weight_prefix = weight_name.removesuffix(".weight")
+            # Transformers fine-grained FP8 uses activation_scale while ModelOpt uses input_scale.
+            # Ref: https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/mistral3.py#L357-L361
+            input_scale_name = next((name for name in (
+                weight_prefix + ".input_scale",
+                weight_prefix + ".activation_scale",
+                weight_prefix + "_activation_scale",
+                weight_prefix + ".qscale_act",
+            ) if name in self.model_tensors), None)
+            input_scale = None
+            if input_scale_name is not None:
+                input_scale = self.model_tensors[input_scale_name]()
+                if prod(input_scale.shape) != 1:
+                    raise ValueError(f"FP8 input scale {input_scale_name!r} must be a scalar")
+                input_scale = LazyTorchTensor.to_eager(input_scale).float().flatten()
+                consumed.append(input_scale_name)
+
+            consumed.append(scale_name)
+
+            expert_match = expert_pattern.search(weight_name)
+            if expert_match:
+                expert_id = int(expert_match.group(2))
+                for new_name in new_names:
+                    target_name = new_name.replace(".weight", ".scale")
+                    entries = scale_tensors.setdefault(target_name, [])
+                    assert isinstance(entries, list)
+                    cast(list[tuple[int, float]], entries).append((expert_id, float(scale[0])))
+                    if input_scale is not None:
+                        target_name = new_name.replace(".weight", ".input_scale")
+                        entries = input_scale_tensors.setdefault(target_name, [])
+                        assert isinstance(entries, list)
+                        cast(list[tuple[int, float]], entries).append((expert_id, float(input_scale[0])))
+            else:
+                for new_name in new_names:
+                    mapped = self.tensor_map.get_type_and_name(new_name, try_suffixes=(".weight",)) if len(shape) == 3 else None
+                    is_packed_expert = mapped is not None and mapped[0] in (
+                        gguf.MODEL_TENSOR.FFN_GATE_EXP, gguf.MODEL_TENSOR.FFN_UP_EXP, gguf.MODEL_TENSOR.FFN_DOWN_EXP,
+                    )
+                    # Packed expert weights need one scale value per expert.
+                    n_scales = shape[0] if is_packed_expert else 1
+                    scale_tensors[new_name.replace(".weight", ".scale")] = np.repeat(scale.numpy(), n_scales)
+                    if input_scale is not None:
+                        input_scale_tensors[new_name.replace(".weight", ".input_scale")] = np.repeat(input_scale.numpy(), n_scales)
+
+        for name in consumed:
+            self.model_tensors.pop(name, None)
+
+        for name, values in chain(scale_tensors.items(), input_scale_tensors.items()):
+            if isinstance(values, list):
+                values.sort(key=lambda item: item[0])
+                scale = np.array([item[1] for item in values], dtype=np.float32)
+            else:
+                scale = values.astype(np.float32)
+            if not np.allclose(scale, 1.0, atol=1e-6):
+                logger.info(f"  + {name} (FP8 scale, shape [{scale.size}])")
+                self.gguf_writer.add_tensor(name, scale)
+
     @classmethod
     def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
         name, gen = item
@@ -625,6 +800,26 @@ class ModelBase:
             raise ValueError(f"Can not map tensor {name!r}")
         return new_name
 
+    def prepare_gate_up_fusion(self) -> None:
+        self._unfusable_gate_up_layers.clear()
+        if not self.fuse_gate_up_exps or self._fp8_as_q8:
+            return
+
+        for name, gen in self.model_tensors.items():
+            merged_name = re.sub(r"(\.experts)\.\d+\.", r"\1.", name)
+            mapped = self.tensor_map.get_type_and_name(merged_name, try_suffixes=(".weight",))
+            if mapped is None:
+                continue
+            tensor_type, new_name = mapped
+            if tensor_type not in (gguf.MODEL_TENSOR.FFN_GATE_EXP, gguf.MODEL_TENSOR.FFN_UP_EXP, gguf.MODEL_TENSOR.FFN_GATE_UP_EXP) or not new_name.endswith(".weight"):
+                continue
+            if gen().dtype == torch.float8_e4m3fn:
+                bid = next(int(part) for part in new_name.split(".") if part.isdecimal())
+                self._unfusable_gate_up_layers.add(bid)
+
+        for bid in sorted(self._unfusable_gate_up_layers):
+            logger.warning(f"Skipping --fuse-gate-up-exps for layer {bid}: fusion of preserved FP8 expert projections is not supported")
+
     def prepare_qkv_fusion(self) -> None:
         self._fusable_qkv_weight_layers.clear()
         self._fusable_qkv_bias_layers.clear()
@@ -638,6 +833,7 @@ class ModelBase:
         }
         weights: dict[int, set[gguf.MODEL_TENSOR]] = {}
         biases: dict[int, set[gguf.MODEL_TENSOR]] = {}
+        fp8_layers: set[int] = set()
 
         for name in self.model_tensors:
             mapped = self.tensor_map.get_type_and_name(name, try_suffixes=(".weight", ".bias"))
@@ -652,12 +848,19 @@ class ModelBase:
                 continue
             if new_name.endswith(".weight"):
                 weights.setdefault(bid, set()).add(tensor_type)
+                if not self._fp8_as_q8 and self.model_tensors[name]().dtype == torch.float8_e4m3fn:
+                    fp8_layers.add(bid)
             elif new_name.endswith(".bias"):
                 biases.setdefault(bid, set()).add(tensor_type)
 
         for bid, weight_types in weights.items():
             bias_types = biases.get(bid, set())
             if weight_types == qkv_types and (not bias_types or bias_types == qkv_types):
+                # NVFP4 weights use separate reordering and repacking hooks and are exported before this fusion pass.
+                # Separate FP8 scales cannot be represented by one fused QKV scale.
+                if bid in fp8_layers:
+                    logger.warning(f"Skipping --fuse-qkv for layer {bid}: fusion of preserved FP8 projections is not supported")
+                    continue
                 self._fusable_qkv_weight_layers.add(bid)
                 if bias_types:
                     self._fusable_qkv_bias_layers.add(bid)
@@ -846,7 +1049,8 @@ class ModelBase:
         new_name = self.map_tensor_name(name)
 
         # Handle gate/up expert tensor fusion if enabled
-        if self.fuse_gate_up_exps and bid is not None:
+        # Preserved FP8 gate/up scales must stay with their separate projections.
+        if self.fuse_gate_up_exps and bid is not None and bid not in self._unfusable_gate_up_layers:
             if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.FFN_GATE_EXP, bid):
                 self._gate_exp_buffer[bid] = data_torch
             elif self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.FFN_UP_EXP, bid):
@@ -903,8 +1107,46 @@ class ModelBase:
 
         return [(new_name, data_torch)]
 
+    def _tensor_requires_f32(self, new_name: str, bid: int | None, n_dims: int) -> bool:
+        # Conditions should closely match those in llama_model_quantize_internal in llama.cpp.
+        return n_dims <= 1 or new_name.endswith("_norm.weight") or (
+            any(
+                self.match_model_tensor_name(new_name, key, bid)
+                for key in (
+                    gguf.MODEL_TENSOR.FFN_GATE_INP,
+                    gguf.MODEL_TENSOR.FFN_GATE_INP_SHEXP,
+                    gguf.MODEL_TENSOR.POS_EMBD,
+                    gguf.MODEL_TENSOR.TOKEN_TYPES,
+                    gguf.MODEL_TENSOR.SSM_CONV1D,
+                    gguf.MODEL_TENSOR.SHORTCONV_CONV,
+                    gguf.MODEL_TENSOR.TIME_MIX_FIRST,
+                    gguf.MODEL_TENSOR.TIME_MIX_W1,
+                    gguf.MODEL_TENSOR.TIME_MIX_W2,
+                    gguf.MODEL_TENSOR.TIME_MIX_DECAY_W1,
+                    gguf.MODEL_TENSOR.TIME_MIX_DECAY_W2,
+                    gguf.MODEL_TENSOR.TIME_MIX_LERP_FUSED,
+                    gguf.MODEL_TENSOR.POSNET_NORM1,
+                    gguf.MODEL_TENSOR.POSNET_NORM2,
+                    gguf.MODEL_TENSOR.V_ENC_EMBD_POS,
+                    gguf.MODEL_TENSOR.A_ENC_EMBD_POS,
+                    gguf.MODEL_TENSOR.ALTUP_CORRECT_COEF,
+                    gguf.MODEL_TENSOR.ALTUP_PREDICT_COEF,
+                    # Kimi KDA conv weights should be F32
+                    gguf.MODEL_TENSOR.SSM_CONV1D_Q,
+                    gguf.MODEL_TENSOR.SSM_CONV1D_K,
+                    gguf.MODEL_TENSOR.SSM_CONV1D_V,
+                    # DSA indexer weights should be F32
+                    gguf.MODEL_TENSOR.INDEXER_PROJ,
+                )
+            )
+            or new_name[-7:] not in (".weight", ".lora_a", ".lora_b")
+        )
+
     def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool:
-        del new_name, bid  # unused
+        if name in self._fp8_e4m3_preserved and n_dims >= 2:
+            if self._tensor_requires_f32(new_name, bid, n_dims):
+                return gguf.GGMLQuantizationType.F32
+            return gguf.GGMLQuantizationType.F8_E4M3
         # Force FP8-original tensors to Q8_0 when requested; Q8_0 is faster than F16/BF16.
         if self._fp8_as_q8 and name in self._fp8_dequantized and n_dims >= 2:
             return gguf.GGMLQuantizationType.Q8_0
@@ -1035,7 +1277,8 @@ class ModelBase:
         expert_scales: dict[tuple[int, str], list[tuple[int, float]]] = {}
         expert_input_scales: dict[tuple[int, str], list[tuple[int, float]]] = {}
         expert_shapes: dict[tuple[int, str], list[int]] = {}
-        n_experts = self.find_hparam(["num_local_experts", "num_experts"], optional=True) or 0
+        expert_names: dict[tuple[int, str], str] = {}
+        n_experts = self.find_hparam(["num_local_experts", "num_experts", "n_routed_experts"], optional=True) or 0
         consumed: list[str] = []
 
         for name in self.model_tensors.keys():
@@ -1046,14 +1289,17 @@ class ModelBase:
             input_scale_name = name.replace(".weight", ".input_scale")
             if scale_name not in self.model_tensors:
                 continue
-            # Force eager materialization of lazy tensors
-            weight = LazyTorchTensor.to_eager(self.model_tensors[name]())
-            scale = LazyTorchTensor.to_eager(self.model_tensors[scale_name]())
+            weight = self.model_tensors[name]()
+            scale = self.model_tensors[scale_name]()
 
-            # Skip non-NVFP4 tensors(e.g. 1D scale, or float8 weight)
-            if scale.ndim < 2 or weight.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
+            # Leave FP8 weights and their scales for _prepare_fp8_e4m3_tensors.
+            if weight.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) or scale.ndim < 2:
                 continue
+            if weight.ndim != 2 or scale.shape[0] != weight.shape[0] or scale.shape[1] * 8 != weight.shape[1]:
+                raise ValueError(f"NVFP4 weight {name!r} has incompatible shapes: weight {list(weight.shape)}, scale {list(scale.shape)}")
 
+            weight = LazyTorchTensor.to_eager(weight)
+            scale = LazyTorchTensor.to_eager(scale)
             scale2 = LazyTorchTensor.to_eager(self.model_tensors.get(scale2_name, lambda: torch.tensor(1.0))())
             input_scale = LazyTorchTensor.to_eager(self.model_tensors.get(input_scale_name, lambda: torch.tensor(1.0))())
 
@@ -1080,6 +1326,7 @@ class ModelBase:
                     expert_scales[key] = []
                     expert_input_scales[key] = []
                     expert_shapes[key] = shape
+                    expert_names[key] = name
                 expert_blocks[key].append((expert_id, raw.copy()))
                 # Collect per-expert scale2 (scalar per expert)
                 expert_scales[key].append((expert_id, float(scale2.float().sum())))
@@ -1088,13 +1335,13 @@ class ModelBase:
 
                 # Flush when all experts for this (layer, proj) are collected
                 if n_experts > 0 and len(expert_blocks[key]) >= n_experts:
-                    self._flush_nvfp4_experts(key, expert_blocks, expert_scales, expert_input_scales, expert_shapes, bid, proj_type)
+                    self._flush_nvfp4_experts(key, expert_blocks, expert_scales, expert_input_scales, expert_shapes, expert_names, bid, proj_type)
             else:
                 self._repack_nvfp4(name, weight, scale, scale2, input_scale)
 
         # Flush any remaining experts (fallback if n_experts was unknown)
         for bid, proj_type in list(expert_blocks.keys()):
-            self._flush_nvfp4_experts((bid, proj_type), expert_blocks, expert_scales, expert_input_scales, expert_shapes, bid, proj_type)
+            self._flush_nvfp4_experts((bid, proj_type), expert_blocks, expert_scales, expert_input_scales, expert_shapes, expert_names, bid, proj_type)
 
         # Remove consumed tensors so get_tensors/modify_tensors won't see them
         for name in consumed:
@@ -1105,7 +1352,7 @@ class ModelBase:
             if name.endswith((".k_scale", ".v_scale")):
                 del self.model_tensors[name]
 
-    def _flush_nvfp4_experts(self, key, expert_blocks, expert_scales, expert_input_scales, expert_shapes, bid, proj_type):
+    def _flush_nvfp4_experts(self, key, expert_blocks, expert_scales, expert_input_scales, expert_shapes, expert_names, bid, proj_type):
         experts = expert_blocks.pop(key)
         scales = expert_scales.pop(key)
         input_scales = expert_input_scales.pop(key)
@@ -1117,7 +1364,7 @@ class ModelBase:
         new_name = self.map_tensor_name(merged_name)
         logger.info(f"Repacked {new_name} with shape [{len(experts)}, {shape[0]}, {shape[1]}] and quantization NVFP4")
         self.gguf_writer.add_tensor(new_name, merged, raw_dtype=gguf.GGMLQuantizationType.NVFP4)
-        self._tag_prec_a4(merged_name, new_name)
+        self._tag_prec_a4(expert_names.pop(key), new_name)
 
         scales.sort(key=lambda x: x[0])
         self._write_scales_tensor(new_name.replace(".weight", ".scale"), [s[1] for s in scales])
@@ -1209,8 +1456,10 @@ class ModelBase:
                             self.model_tensors[input_scale_name] = inverse_scale(self.model_tensors.pop(name))
             self._generate_nvfp4_tensors()
 
+        self._prepare_fp8_e4m3_tensors()
         self.dequant_model()
 
+        self.prepare_gate_up_fusion()
         self.prepare_qkv_fusion()
 
         # Handle empty tensor_map for models with block_count=0 (like MobileNetV5)
@@ -1224,10 +1473,21 @@ class ModelBase:
             if name.endswith((".attention.masked_bias", ".attention.bias", ".rotary_emb.inv_freq")):
                 continue
 
+            # FP8 transforms must only reorder values; norms and vectors use numeric F32 transforms.
+            preserve_fp8_e4m3 = self._preserve_fp8 and data_torch.dtype == torch.float8_e4m3fn and data_torch.ndim >= 2 and not name.endswith(("norm.weight", "norm_f.weight"))
+            if preserve_fp8_e4m3:
+                for new_name in self._map_fp8_weight_names(name):
+                    mapped_bid = next((int(part) for part in new_name.split(".") if part.isdecimal()), None)
+                    if self._tensor_requires_f32(new_name, mapped_bid, data_torch.ndim):
+                        preserve_fp8_e4m3 = False
+                        break
+                if preserve_fp8_e4m3:
+                    self._fp8_e4m3_preserved.add(name)
+
             old_dtype = data_torch.dtype
 
             # convert any unsupported data types to float32
-            if data_torch.dtype not in (torch.float16, torch.float32):
+            if data_torch.dtype not in (torch.float16, torch.float32) and name not in self._fp8_e4m3_preserved:
                 data_torch = data_torch.to(torch.float32)
 
             # use the first number-like part of the tensor name as the block id
@@ -1238,51 +1498,14 @@ class ModelBase:
                     break
 
             for new_name, data_torch in (self.modify_tensors(data_torch, name, bid)):
-                # TODO: why do we squeeze here?
-                # data = data_torch.squeeze().numpy()
-                data = data_torch.numpy()
-
-                n_dims = len(data.shape)
+                n_dims = len(data_torch.shape)
                 data_qtype: gguf.GGMLQuantizationType | bool = self.tensor_force_quant(name, new_name, bid, n_dims)
 
                 # Most of the codebase that takes in 1D tensors or norms only handles F32 tensors
                 if n_dims <= 1 or new_name.endswith("_norm.weight"):
                     data_qtype = gguf.GGMLQuantizationType.F32
 
-                # Conditions should closely match those in llama_model_quantize_internal in llama.cpp
-                # Some tensor types are always in float32
-                if data_qtype is False and (
-                    any(
-                        self.match_model_tensor_name(new_name, key, bid)
-                        for key in (
-                            gguf.MODEL_TENSOR.FFN_GATE_INP,
-                            gguf.MODEL_TENSOR.FFN_GATE_INP_SHEXP,
-                            gguf.MODEL_TENSOR.POS_EMBD,
-                            gguf.MODEL_TENSOR.TOKEN_TYPES,
-                            gguf.MODEL_TENSOR.SSM_CONV1D,
-                            gguf.MODEL_TENSOR.SHORTCONV_CONV,
-                            gguf.MODEL_TENSOR.TIME_MIX_FIRST,
-                            gguf.MODEL_TENSOR.TIME_MIX_W1,
-                            gguf.MODEL_TENSOR.TIME_MIX_W2,
-                            gguf.MODEL_TENSOR.TIME_MIX_DECAY_W1,
-                            gguf.MODEL_TENSOR.TIME_MIX_DECAY_W2,
-                            gguf.MODEL_TENSOR.TIME_MIX_LERP_FUSED,
-                            gguf.MODEL_TENSOR.POSNET_NORM1,
-                            gguf.MODEL_TENSOR.POSNET_NORM2,
-                            gguf.MODEL_TENSOR.V_ENC_EMBD_POS,
-                            gguf.MODEL_TENSOR.A_ENC_EMBD_POS,
-                            gguf.MODEL_TENSOR.ALTUP_CORRECT_COEF,
-                            gguf.MODEL_TENSOR.ALTUP_PREDICT_COEF,
-                            # Kimi KDA conv weights should be F32
-                            gguf.MODEL_TENSOR.SSM_CONV1D_Q,
-                            gguf.MODEL_TENSOR.SSM_CONV1D_K,
-                            gguf.MODEL_TENSOR.SSM_CONV1D_V,
-                            # DSA indexer weights should be F32
-                            gguf.MODEL_TENSOR.INDEXER_PROJ,
-                        )
-                    )
-                    or new_name[-7:] not in (".weight", ".lora_a", ".lora_b")
-                ):
+                if data_qtype is False and self._tensor_requires_f32(new_name, bid, n_dims):
                     data_qtype = gguf.GGMLQuantizationType.F32
 
                 if data_qtype is False and any(
@@ -1320,16 +1543,22 @@ class ModelBase:
                     else:
                         raise ValueError(f"Unknown file type: {self.ftype.name}")
 
+                if data_torch.dtype == torch.float8_e4m3fn:
+                    data = data_torch.view(torch.uint8).numpy() if data_qtype == gguf.GGMLQuantizationType.F8_E4M3 else data_torch.float().numpy()
+                else:
+                    data = data_torch.numpy()
+
                 # a chunked tensor quantizes as one chunk at a time, while it is written
                 quantize = data.quantize if isinstance(data, gguf.LazyChunkedTensor) else (
                     lambda qtype, d=data: gguf.quants.quantize(d, qtype))
 
-                try:
-                    data = quantize(data_qtype)
-                except gguf.QuantError as e:
-                    logger.warning("%s, %s", e, "falling back to F16")
-                    data_qtype = gguf.GGMLQuantizationType.F16
-                    data = quantize(data_qtype)
+                if not (data_qtype == gguf.GGMLQuantizationType.F8_E4M3 and data.dtype == np.uint8):
+                    try:
+                        data = quantize(data_qtype)
+                    except gguf.QuantError as e:
+                        logger.warning("%s, %s", e, "falling back to F16")
+                        data_qtype = gguf.GGMLQuantizationType.F16
+                        data = quantize(data_qtype)
 
                 shape = gguf.quant_shape_from_byte_shape(data.shape, data_qtype) if data.dtype == np.uint8 else data.shape
 
@@ -1370,6 +1599,8 @@ class ModelBase:
                 self.ftype = gguf.LlamaFileType.MOSTLY_NVFP4
             elif self._is_mxfp4:
                 self.ftype = gguf.LlamaFileType.MOSTLY_MXFP4_MOE
+            elif self._fp8_e4m3_preserved:
+                self.ftype = gguf.LlamaFileType.MOSTLY_F8_E4M3
 
         # Generate parameter weight class (useful for leader boards) if not yet determined
         if self.metadata.size_label is None and total_params > 0:
@@ -2672,7 +2903,13 @@ class TextModel(ModelBase):
         )
 
         if vocab.tokenizer_type == MistralTokenizerType.tekken:
-            self.gguf_writer.add_tokenizer_pre("tekken")
+            with open(vocab.fname_tokenizer, encoding="utf-8") as f:
+                pattern = json.load(f)["config"]["pattern"]
+            # newer tekken.json split numbers into right-aligned groups of 3 digits instead of single digits
+            if "\\p{N}{1,3}(?=(?:\\p{N}{3})*(?!\\p{N}))" in pattern:
+                self.gguf_writer.add_tokenizer_pre("tekken-num3")
+            else:
+                self.gguf_writer.add_tokenizer_pre("tekken")
             self.gguf_writer.add_token_merges(
                 vocab.extract_vocab_merges_from_model()
             )
@@ -3149,7 +3386,7 @@ def get_model_architecture(hparams: dict[str, Any], model_type: ModelType) -> st
     # Kimi-K3's text_config reports "KimiLinearForCausalLM", which is the older
     # Kimi-Linear-48B architecture and cannot load K3 (no attention residuals,
     # latent MoE, situ, ...). Route on the top-level architecture instead.
-    if model_type == ModelType.TEXT and arch in ("NEOChatModel", "StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration", "KimiK3ForConditionalGeneration"):
+    if model_type == ModelType.TEXT and arch in ("NEOChatModel", "StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration", "KimiK3ForConditionalGeneration", "Nemotron3_5TranscribeForConditionalGeneration"):
         return arch
 
     # if "architectures" is found in the sub-config, use that instead

conversion/clover.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+from typing import Callable, Iterable, TYPE_CHECKING
+
+from .base import ModelBase, gguf
+from .qwen import _LinearAttentionVReorderBase
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+
+@ModelBase.register("CloverMoeModelLM")
+@ModelBase.example("Diffuspace/Clover-35B-A3B-Preview")
+class CloverMoeModel(_LinearAttentionVReorderBase):
+    model_arch = gguf.MODEL_ARCH.CLOVER_MOE
+    supports_mtp_export = False
+
+    def __init__(self, *args, **kwargs):
+        self.no_mtp = True
+        super().__init__(*args, **kwargs)
+
+    def set_vocab(self):
+        self._set_vocab_gpt2()
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        self.gguf_writer.add_expert_shared_count(self.hparams.get("num_shared_experts", 1))
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, _ = item
+        if name.endswith((".conv_right_weight", ".bd_softmax_logit_scale", ".bd_softmax_history_bias", ".bd_softmax_history_c")):
+            return None
+
+        return super().filter_tensors(item)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name.endswith(".A_log") or (name.endswith("norm.weight") and not name.endswith("linear_attn.norm.weight")):
+            data_torch = data_torch.float()
+
+        yield from super().modify_tensors(data_torch, name, bid)

conversion/deepseek.py

@@ -236,6 +236,12 @@ class DeepseekModel(TextModel):
 class DeepseekV2Model(TextModel):
     model_arch = gguf.MODEL_ARCH.DEEPSEEK2
 
+    def _map_fp8_weight_names(self, name: str) -> tuple[str, ...]:
+        # MLA projections do not have runtime weight-scale support.
+        if name.endswith((".q_a_proj.weight", ".q_b_proj.weight", ".kv_a_proj_with_mqa.weight", ".kv_b_proj.weight")):
+            return ()
+        return super()._map_fp8_weight_names(name)
+
     # TODO @ngxson : remove this when we support MTP for deepseek models
     skip_mtp = True
 

conversion/lfm3.py

@@ -453,9 +453,6 @@ class LFM3AudioMmprojModel(MmprojModel):
     def tensor_force_quant(self, name, new_name, bid, n_dims):
         if new_name.startswith("a.conv1d.") or "conv_dw" in new_name or new_name.startswith("a.gen.wav.quant."):
             return gguf.GGMLQuantizationType.F32
-        # the codec decoder is small and its log-magnitude/phase output does not survive quantization
-        if new_name.startswith("a.gen.wav."):
-            return gguf.GGMLQuantizationType.F16
         return super().tensor_force_quant(name, new_name, bid, n_dims)
 
     def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:

conversion/limite.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+from typing import Iterable, TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import ModelBase, TextModel, gguf
+
+
+@ModelBase.register("LimiteForCausalLM")
+@ModelBase.example("paradigma-inc/limite-1b-violetto")
+class LimiteModel(TextModel):
+    model_arch = gguf.MODEL_ARCH.LIMITE
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        hp = self.hparams
+
+        expected = {
+            "pos_mode": "rope",
+            "qk_norm": "rms_pre_rope",
+            "rope_style": "interleaved_pairs_odd_lane_sign_flip",
+            "rope_per_layer": False,
+            "rms_norm_has_weight": False,
+            "rms_norm_eps_mode": "torch_finfo_default",
+            "mlp_type": "swiglu",
+            "attn_gate_applied": "per_head_before_o_proj",
+            "ve_head_slice": "first_num_key_value_heads",
+            "ve_applied_before_qk_norm": True,
+            "global_nope": True,
+            "global_window": -1,
+            "mudd_r_site": "resid",
+            "tie_word_embeddings": True,
+            "final_softcap": 0.0,
+        }
+        for key, value in expected.items():
+            if hp.get(key) != value:
+                raise NotImplementedError(f"Limite: unsupported {key}={hp.get(key)!r}, expected {value!r}")
+        if hp["softcap_logits"]["kind"] != "sigmoid":
+            raise NotImplementedError(f"Limite: unsupported softcap_logits kind {hp['softcap_logits']['kind']!r}")
+
+        self.n_head    = hp["num_attention_heads"]
+        self.n_head_kv = hp["num_key_value_heads"]
+        self.head_dim  = hp["head_dim"]
+        self.n_pairs   = hp["rope_n_pairs"]
+
+        self.mudd_taps: dict[int, list[int]] = {}
+        if hp.get("mudd"):
+            if not hp.get("mudd_mlp"):
+                raise NotImplementedError("Limite: mudd without mudd_mlp is not supported")
+            self.mudd_taps = {int(il): [int(t) for t in taps] for il, taps in hp["mudd_tap_idx"].items()}
+            assert all(len(taps) == hp["mudd_taps"] for taps in self.mudd_taps.values())
+
+    def set_vocab(self):
+        self._set_vocab_gpt2()
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        hp = self.hparams
+
+        self.gguf_writer.add_key_length(self.head_dim)
+        self.gguf_writer.add_value_length(self.head_dim)
+
+        # torch rms_norm with eps=None uses the fp32 epsilon, even for bf16 inputs
+        self.gguf_writer.add_layer_norm_rms_eps(torch.finfo(torch.float32).eps)
+        self.gguf_writer.add_attention_scale(hp["attention_softmax_scale"])
+        self.gguf_writer.add_xsa_eps(hp["xsa_normalize_eps"] if hp.get("xsa") else 0.0)
+
+        # HF freq_i = base^(-i/(n_pairs-1)), ggml uses freq_base^(-i/n_pairs)
+        self.gguf_writer.add_rope_dimension_count(2 * self.n_pairs)
+        self.gguf_writer.add_rope_freq_base(hp["rope_base_local"] ** (self.n_pairs / (self.n_pairs - 1)))
+
+        # HF keeps keys k >= q - window, so a window of W spans W + 1 tokens
+        global_layers = set(hp["global_layers"])
+        self.gguf_writer.add_sliding_window(hp["sliding_window"] + 1)
+        self.gguf_writer.add_sliding_window_pattern([il not in global_layers for il in range(self.block_count)])
+
+        cap = hp["softcap_logits"]
+        self.gguf_writer.add_final_logit_sigmoid_capping(cap["a"], cap["b"], cap["c"])
+
+        if self.mudd_taps:
+            n_taps = hp["mudd_taps"]
+            tap_indices = [-1] * (self.block_count * n_taps)
+            for il, taps in self.mudd_taps.items():
+                tap_indices[il * n_taps:(il + 1) * n_taps] = taps
+            self.gguf_writer.add_mudd_feed_forward_length(hp["mudd_inter"])
+            self.gguf_writer.add_mudd_tap_count(n_taps)
+            self.gguf_writer.add_mudd_tap_indices(tap_indices)
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        # tiny mixing weights, kept at full precision as in the reference
+        # the mudd tensors come from a global source tensor, so take the block id from the new name
+        if new_name.startswith("blk."):
+            bid = int(new_name.split(".")[1])
+        for key in (gguf.MODEL_TENSOR.ATTN_GATE, gguf.MODEL_TENSOR.ATTN_VE_GATE, gguf.MODEL_TENSOR.MUDD_DOWN, gguf.MODEL_TENSOR.MUDD_UP):
+            if self.match_model_tensor_name(new_name, key, bid):
+                return gguf.GGMLQuantizationType.F32
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
+
+    def _swap_rope_pairs(self, w: Tensor, n_head: int) -> Tensor:
+        # HF rotates each pair by -theta; swapping the pair lanes gives the ggml +theta rotation.
+        # q and k are both swapped, so the attention scores do not change.
+        idx = torch.arange(self.head_dim)
+        idx[:2 * self.n_pairs] = idx[:2 * self.n_pairs].view(-1, 2).flip(-1).reshape(-1)
+        w = w.reshape(n_head, self.head_dim, -1)
+        return w[:, idx, :].reshape(n_head * self.head_dim, -1)
+
+    def _pad_cols(self, w: Tensor) -> Tensor:
+        # gates read only the first channels of the input, pad them to n_embd
+        return torch.nn.functional.pad(w, (0, self.hparams["hidden_size"] - w.shape[-1]))
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        hp = self.hparams
+        T = gguf.MODEL_TENSOR
+
+        def load(n: str) -> Tensor:
+            return self.model_tensors[n]().float()
+
+        if name == "model.value_embeds.weight":
+            ve = data_torch.reshape(data_torch.shape[0], -1, hp["ve_dim"])[:, :self.n_head_kv]
+            ve = torch.nn.functional.pad(ve, (0, self.head_dim - hp["ve_dim"]))
+            yield self.format_tensor_name(T.VALUE_EMBD), ve.reshape(ve.shape[0], -1) * hp["ve_gate_scale"]
+            return
+
+        if name == "model.mudd.dense1":
+            for il in self.mudd_taps:
+                yield self.format_tensor_name(T.MUDD_DOWN, il), data_torch
+            return
+        if name == "model.mudd.dense2":
+            # rows [0, n) mix the attention input, rows [n, 2n) mix the residual
+            dense2_mlp = load("model.mudd.dense2_mlp")
+            bias       = load("model.mudd.bias")
+            bias_mlp   = load("model.mudd.bias_mlp")
+            for il, taps in self.mudd_taps.items():
+                n = len(taps)
+                w = torch.cat([data_torch[il, :n], dense2_mlp[il, :n]], 0)
+                b = torch.cat([bias[il, :n], bias_mlp[il, :n]], 0)
+                yield self.format_tensor_name(T.MUDD_UP, il), w
+                yield self.format_tensor_name(T.MUDD_UP, il, ".bias"), b
+            return
+        if name.startswith("model.mudd."):
+            return
+
+        if bid is None:
+            yield from super().modify_tensors(data_torch, name, bid)
+            return
+
+        prefix = f"model.layers.{bid}."
+        name = name.removeprefix(prefix)
+
+        if name in ("self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.v_proj.weight"):
+            data_torch = data_torch * load(prefix + "self_attn.qkv_scale")
+            if name == "self_attn.q_proj.weight":
+                data_torch = self._swap_rope_pairs(data_torch, self.n_head)
+            elif name == "self_attn.k_proj.weight":
+                data_torch = self._swap_rope_pairs(data_torch, self.n_head_kv)
+        elif name == "self_attn.o_proj.weight":
+            data_torch = data_torch * load(prefix + "self_attn.o_scale") * load(prefix + "post_lambda_attn") * hp["attn_gate_scale"]
+        elif name == "mlp.down_proj.weight":
+            data_torch = data_torch * load(prefix + "post_lambda_mlp")
+        elif name == "self_attn.attn_gate":
+            yield self.format_tensor_name(T.ATTN_GATE, bid), self._pad_cols(data_torch)
+            return
+        elif name == "self_attn.ve_gate":
+            yield self.format_tensor_name(T.ATTN_VE_GATE, bid), self._pad_cols(data_torch[:self.n_head_kv])
+            return
+        elif name == "self_attn.xsa_alpha":
+            if hp.get("xsa") and bid in hp["xsa_layers"]:
+                yield self.format_tensor_name(T.ATTN_XSA_ALPHA, bid), torch.tanh(data_torch)
+            return
+        elif name == "resid_lambda_attn":
+            yield self.format_tensor_name(T.ATTN_RESID_SCALE, bid), data_torch.reshape(1)
+            return
+        elif name == "resid_lambda_mlp":
+            yield self.format_tensor_name(T.FFN_RESID_SCALE, bid), data_torch.reshape(1)
+            return
+        elif name in ("self_attn.qkv_scale", "self_attn.o_scale", "post_lambda_attn", "post_lambda_mlp"):
+            # folded into the projections above
+            return
+
+        yield from super().modify_tensors(data_torch, prefix + name, bid)

conversion/mistral4.py

@@ -0,0 +1,217 @@
+from __future__ import annotations
+
+from typing import Callable, Iterable, TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import ModelBase, gguf
+
+from .mistral import MistralModel, MistralMoeModel
+
+
+@ModelBase.register("MistralLarge4ForConditionalGeneration")
+class MistralLarge4Model(MistralMoeModel):
+    model_arch = gguf.MODEL_ARCH.MISTRAL_LARGE4
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        n_layer = self.hparams["num_hidden_layers"]
+        dsa = self.hparams["dsa"]
+
+        def expand(pattern: str) -> list[str]:
+            values = pattern.split(",")
+            return [values[i % len(values)] for i in range(n_layer)]
+
+        self.layer_has_rope = [v == "rope" for v in expand(self.hparams.get("nope") or "rope")]
+        self.layer_indexer_full = [v == "full" for v in expand(dsa.get("sharing") or "full")]
+
+        kv_sharing = self.hparams.get("kv_sharing") or {}
+        if "source_layers_str" in kv_sharing:
+            self.layer_kv_source = [int(v) for v in kv_sharing["source_layers_str"].split(",")]
+        else:
+            assert not kv_sharing, f"unsupported kv_sharing config: {kv_sharing}"
+            self.layer_kv_source = list(range(n_layer))
+        assert len(self.layer_kv_source) == n_layer
+
+        # note: the rope dim of the indexer must match the main attention
+        assert dsa.get("rope_head_dim", -1) in (-1, self.hparams["qk_rope_head_dim"])
+        assert dsa.get("match_nope", False), "indexer must follow the NoPE pattern of the main attention"
+        assert dsa.get("ssmax_scaling", False), "indexer without SSMax is not supported"
+        assert self.hparams.get("q_lora_rank") is None
+        assert not self.hparams["yarn"].get("apply_scale", True), "only YaRN without magnitude scaling is supported"
+
+    def dequant_model(self):
+        quant_config = self.hparams.get("quantization")
+        if quant_config is not None:
+            assert quant_config["qformat_weight"] == "fp8_e4m3"
+            assert quant_config.get("qscheme_weight") == "BLOCK"
+            self.hparams["quantization_config"] = {
+                "quant_method": "fp8",
+                "weight_block_size": [quant_config["block_size_out"], quant_config["block_size_in"]],
+            }
+        return super().dequant_model()
+
+    def set_gguf_parameters(self):
+        # skip the MistralMoeModel overrides, they apply the deepseek2 YaRN scaling
+        super(MistralMoeModel, self).set_gguf_parameters()
+        MistralModel.set_mistral_config(self.gguf_writer, self.hparams)
+
+        dsa = self.hparams["dsa"]
+        self.gguf_writer.add_indexer_head_count(dsa["n_heads"])
+        self.gguf_writer.add_indexer_key_length(dsa["head_dim"])
+        self.gguf_writer.add_indexer_top_k(dsa["topk"])
+        self.gguf_writer.add_indexer_types(self.layer_indexer_full)
+        self.gguf_writer.add_layer_norm_eps(dsa.get("norm_eps", 1e-6))
+
+        self.gguf_writer.add_rope_pattern(self.layer_has_rope)
+        self.gguf_writer.add_kv_source_map(self.layer_kv_source)
+        self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SOFTMAX)
+        self.gguf_writer.add_final_logit_softcapping(self.hparams["softmax_tanh_gating"])
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+
+        # exported separately with --dspark
+        if name.startswith("mtp_"):
+            return None
+
+        # experts need HF names for the merging logic in DeepseekV2Model
+        # note: bare "layers.N.experts" comes from the NVFP4 quantized_layers config
+        if ".experts." in name or name.endswith(".experts"):
+            name = "model." + name.replace(".experts", ".mlp.experts")
+            name = name.replace(".w1.", ".gate_proj.").replace(".w2.", ".down_proj.").replace(".w3.", ".up_proj.")
+
+        # skip the MistralMoeModel renaming, the qscale names are needed for block FP8 dequant
+        return super(MistralMoeModel, cls).filter_tensors((name, gen))
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name == "lerope.alpha":
+            # LeRoPE scales inv_freq by exp(alpha), ggml divides theta by the freq factors
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), torch.exp(-data_torch.float()))
+            return
+
+        if bid is not None and name.endswith((".ssmax_scaling.s", ".ssmax_scaling.b")):
+            tensor = gguf.MODEL_TENSOR.INDEXER_SSMAX if ".indexer." in name else gguf.MODEL_TENSOR.ATTN_SSMAX
+            suffix = ".weight" if name.endswith(".s") else ".bias"
+            yield (self.format_tensor_name(tensor, bid, suffix), data_torch.float())
+            return
+
+        if bid is not None and ".indexer.q_input_norm." in name:
+            suffix = ".weight" if name.endswith(".weight") else ".bias"
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_Q_NORM, bid, suffix), data_torch)
+            return
+
+        if bid is not None and ".indexer." in name:
+            name = "model." + name.replace(".attention.", ".self_attn.")
+
+        # NoPE layers have a larger non-rotary head, so wkv_b is split per layer
+        if bid is not None and name.endswith(".wkv_b.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"]
+            if not self.layer_has_rope[bid]:
+                qk_nope_head_dim += self.hparams["qk_rope_head_dim"]
+
+            assert data_torch.shape[0] == n_head * (qk_nope_head_dim + v_head_dim)
+
+            kv_b = data_torch.view(n_head, qk_nope_head_dim + v_head_dim, data_torch.shape[-1])
+            k_b, v_b = torch.split(kv_b, [qk_nope_head_dim, v_head_dim], dim=1)
+            k_b = k_b.transpose(1, 2)
+
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K_B, bid), k_b)
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V_B, bid), v_b)
+            return
+
+        yield from super().modify_tensors(data_torch, name, bid)
+
+
+class MistralLarge4DSparkModel(MistralLarge4Model):
+    """Mistral Large 4 DSpark draft (mtp_dspark.*) as a DFLASH GGUF, export with --dspark"""
+    model_arch = gguf.MODEL_ARCH.DFLASH
+
+    _DSPARK_ROOT: dict[str, gguf.MODEL_TENSOR] = {
+        "backbone.target_projection.weight": gguf.MODEL_TENSOR.FC,
+        "backbone.target_norm.weight":       gguf.MODEL_TENSOR.ENC_OUTPUT_NORM,
+        "backbone.final_norm.weight":        gguf.MODEL_TENSOR.OUTPUT_NORM,
+        "backbone.mask_embedding":           gguf.MODEL_TENSOR.DSPARK_MASK_EMBD,
+        "autoregressive_head.embedding.weight": gguf.MODEL_TENSOR.DSPARK_MARKOV_W1,
+        "autoregressive_head.output.weight":    gguf.MODEL_TENSOR.DSPARK_MARKOV_W2,
+        "autoregressive_head.joint_proj.weight": gguf.MODEL_TENSOR.DSPARK_JOINT_PROJ,
+        "autoregressive_head.joint_proj.bias":   gguf.MODEL_TENSOR.DSPARK_JOINT_PROJ,
+    }
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        mtp = self.hparams.get("mtp") or {}
+        dspark = mtp.get("dspark")
+        if dspark is None:
+            raise ValueError("Mistral Large 4 DSpark export needs the mtp.dspark config")
+        assert dspark.get("head_type") == "rnn", f"unsupported DSpark head type: {dspark.get('head_type')}"
+        assert dspark.get("draft_input") == "position_specific", f"unsupported DSpark draft input: {dspark.get('draft_input')}"
+
+        windows = {int(v) for v in (mtp.get("ragged_attention") or "").split(",") if v.strip()}
+        assert len(windows) == 1, f"DSpark draft layers must share one sliding window, got {mtp.get('ragged_attention')}"
+
+        self._n_stage = int(mtp["n_layers"])
+        self._depth = int(mtp["depth"])
+        self._window = windows.pop()
+        self._target_layers = [int(v) for v in mtp["target_activation_layers"].split(",")]
+
+        self.layer_has_rope = [True] * self._n_stage
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+        if not name.startswith("mtp_dspark."):
+            return None
+        name = name.removeprefix("mtp_dspark.")
+        if name.startswith("backbone.layers."):
+            return name.removeprefix("backbone."), gen
+        return f"dspark.{name}", gen
+
+    def set_vocab(self):
+        super().set_vocab()
+        # block uses mask_embd, mask token id is not used
+        self.gguf_writer.add_mask_token_id(0)
+
+    def set_gguf_parameters(self):
+        hp = self.hparams
+        w = self.gguf_writer
+        w.add_block_count(self._n_stage)
+        w.add_context_length(hp["max_position_embeddings"])
+        w.add_embedding_length(hp["hidden_size"])
+        w.add_feed_forward_length(hp["intermediate_size"])
+        w.add_head_count(hp["num_attention_heads"])
+        w.add_head_count_kv(1)
+        w.add_key_length(hp["kv_lora_rank"] + hp["qk_rope_head_dim"])
+        w.add_value_length(hp["kv_lora_rank"])
+        w.add_key_length_mla(hp["qk_nope_head_dim"] + hp["qk_rope_head_dim"])
+        w.add_value_length_mla(hp["v_head_dim"])
+        w.add_kv_lora_rank(hp["kv_lora_rank"])
+        w.add_rope_dimension_count(hp["qk_rope_head_dim"])
+        w.add_rope_freq_base(float(hp["rope_theta"]))
+        w.add_layer_norm_rms_eps(hp["rms_norm_eps"])
+        MistralModel.set_mistral_config(w, hp)
+        w.add_sliding_window(self._window)
+        w.add_final_logit_softcapping(hp["softmax_tanh_gating"])
+        w.add_block_size(self._depth)
+        # draft reads layer output = input of next layer
+        w.add_target_layers([layer + 1 for layer in self._target_layers])
+        w.add_sample_from_anchor(True)
+        w.add_has_confidence_head(False)
+        w.add_file_type(self.ftype)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name.startswith("dspark."):
+            key = name.removeprefix("dspark.")
+            if key not in self._DSPARK_ROOT:
+                raise ValueError(f"Unexpected DSpark tensor {name!r}")
+            suffix = ".bias" if key.endswith(".bias") else ".weight"
+            yield (self.format_tensor_name(self._DSPARK_ROOT[key], None, suffix), data_torch)
+            return
+
+        yield from super().modify_tensors(data_torch, name, bid)

conversion/nemotron.py

@@ -428,11 +428,19 @@ class NemotronHModel(GraniteHybridModel):
     is_moe: bool = False
     supports_mtp_export = True
     _experts: list[dict[str, Tensor]] | None = None
+    # tokens past the embedding rows (left out of the vocab), and special tokens that are printed as text
+    _vocab_skip: tuple[str, ...] = ()
+    _user_defined_tokens: re.Pattern[str] | None = None
 
     _SSM_LAYER_TYPES = {"mamba", "linear_attention"}
     _ATTN_LAYER_TYPES = {"attention", "full_attention"}
     _MLP_LAYER_TYPES = {"moe"}
 
+    def _map_fp8_weight_names(self, name: str) -> tuple[str, ...]:
+        if ".mixer.experts." in name:
+            name = name.replace("backbone.", "model.").replace(".mixer.experts.", ".mlp.experts.")
+        return super()._map_fp8_weight_names(name)
+
     def __init__(self, *args, **kwargs):
         # We have to determine the correct model architecture (MoE vs non-MoE) before
         # calling the parent __init__. This is because the parent constructor
@@ -611,7 +619,11 @@ class NemotronHModel(GraniteHybridModel):
         vocab_size = -(vocab_size // -pad_vocab) * pad_vocab
         self.hparams["vocab_size"] = vocab_size
 
-        assert max(tokenizer.vocab.values()) < vocab_size  # ty: ignore[unresolved-attribute]
+        # draft-model converters (DFlash, LiLiCorr) call this with their own instance, which has neither hook
+        vocab_skip = getattr(self, "_vocab_skip", ())
+        user_defined_tokens = getattr(self, "_user_defined_tokens", None)
+
+        assert max(id_ for tok, id_ in tokenizer.vocab.items() if tok not in vocab_skip) < vocab_size  # ty: ignore[unresolved-attribute]
 
         tokpre = self.get_vocab_base_pre(tokenizer)
 
@@ -633,7 +645,8 @@ class NemotronHModel(GraniteHybridModel):
                         if previous_token != token:
                             logger.info(f"{repr(previous_token)} is encoded and decoded back to {repr(token)} using AutoTokenizer")
 
-                    if added_tokens_decoder[i].special or self.does_token_look_special(token):
+                    printed = user_defined_tokens is not None and user_defined_tokens.fullmatch(token)
+                    if (added_tokens_decoder[i].special or self.does_token_look_special(token)) and not printed:
                         toktypes.append(gguf.TokenType.CONTROL)
                     else:
                         token = token.replace(b"\xe2\x96\x81".decode("utf-8"), " ")  # pre-normalize user-defined spaces
@@ -829,3 +842,140 @@ class NemotronHPuzzleModel(NemotronHModel):
         if item[0].startswith("mtp."):
             return None
         return super().filter_tensors(item)
+
+
+# NeMo speechlm2 SALM checkpoint: a Nemotron-H decoder under llm.*, the speech encoder under perception.*
+@ModelBase.register("Nemotron3_5TranscribeForConditionalGeneration")
+@ModelBase.example("nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B")
+class NemotronTranscribeModel(NemotronHModel):
+    model_arch = gguf.MODEL_ARCH.NEMOTRON_H
+    # the audio placeholder (id vocab_size) has no embedding row; speaker, event and style tags print with the transcript
+    _vocab_skip = ("<|audio|>",)
+    _user_defined_tokens = re.compile(r"<spk:\d+>|<ev:[a-z_]+>|<sty:(?:stt|end):[a-z_]+>")
+    _renames = {
+        "llm.model.embed_tokens.weight": "backbone.embeddings.weight",
+        "llm.model.norm.weight":         "backbone.norm_f.weight",
+        "llm.lm_head.weight":            "lm_head.weight",
+    }
+    # grouped experts: [n_expert, n_in, n_out] -> [n_expert, n_out, n_in]
+    _grouped_experts = {"gate_and_up_projs": "up_proj", "down_projs": "down_proj"}
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+        if name.startswith(("perception.", "ctc_timestamp.")) or "_extra_state" in name:
+            return None
+        if not name.startswith("llm."):
+            # the FP8 and NVFP4 releases store the decoder as backbone.*
+            raise ValueError(f"unexpected tensor {name}: only the BF16 checkpoint layout (llm.*) is supported")
+        if name in cls._renames:
+            name = cls._renames[name]
+        elif name.startswith("llm.model.layers."):
+            name = "backbone.layers." + name[len("llm.model.layers."):]
+        elif name.startswith("llm.mtp.layers."):
+            name = "mtp.layers." + name[len("llm.mtp.layers."):]
+        else:
+            raise ValueError(f"unexpected decoder tensor {name}")
+        return super().filter_tensors((name.replace("._fp32_params.", "."), gen))
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        leaf = name.rsplit(".", 1)[-1]
+        if ".mixer.experts." in name and leaf in self._grouped_experts:
+            if name.startswith("mtp."):
+                assert self._mtp_bid is not None, f"{name}: MTP head is not exported"
+                bid = self._mtp_bid
+            assert bid is not None
+            data_torch = data_torch.permute(0, 2, 1).contiguous()
+            name = f"model.layers.{bid}.mlp.experts.{self._grouped_experts[leaf]}.weight"
+            yield from ModelBase.modify_tensors(self, data_torch, name, bid)
+            return
+        yield from super().modify_tensors(data_torch, name, bid)
+
+
+@ModelBase.register("Nemotron3_5TranscribeForConditionalGeneration")
+class NemotronTranscribeAudioModel(MmprojModel):
+    # NeMo ParallelExpertEncoder: ASR and speaker transformers, speaker activity fused into the ASR states
+    has_vision_encoder = False
+    has_audio_encoder = True
+
+    def get_audio_config(self) -> dict[str, Any] | None:
+        asr = self.global_config["pe_encoder_config"]["asr_encoder_cfg"]
+        return {
+            "hidden_size":         asr["d_model"],
+            "intermediate_size":   int(float(asr.get("ff_expansion", 4)) * asr["d_model"]),
+            "num_hidden_layers":   asr["n_layers"],
+            "num_attention_heads": asr["n_heads"],
+        }
+
+    @staticmethod
+    def _rope_dims(cfg: dict[str, Any]) -> int:
+        assert cfg.get("self_attention_model") == "rope" and cfg.get("subsampling", "feature_stacking") == "feature_stacking"
+        assert cfg.get("attn_mode", "full") == "full" and not cfg.get("xscaling", False) and cfg.get("pre_block_norm", True)
+        assert float(cfg.get("rope_base", 10000.0)) == 10000.0
+        n_rot = int(cfg["d_model"] // cfg["n_heads"] * float(cfg.get("rotary_fraction", 1.0)))
+        assert n_rot >= 2 and n_rot % 2 == 0
+        return n_rot
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        pe = self.global_config["pe_encoder_config"]
+        asr = pe["asr_encoder_cfg"]
+        diar = pe["diarization_model_cfg"]
+        spk = diar["encoder"]
+        assert pe.get("asr_encoder_type") == "transformer" and pe.get("speaker_feature_mode", "thresholded") == "thresholded"
+        assert pe.get("asr_normalize_type") == "per_feature" and pe.get("diar_normalize_type") in (None, "NA")
+        assert asr["subsampling_factor"] == spk["subsampling_factor"] == 8 and asr["feat_in"] == spk["feat_in"]
+        assert self._rope_dims(spk) == spk["d_model"] // spk["n_heads"] and not spk.get("qkv_bias", False)
+        assert diar.get("high_resolution", False)
+        # constants in clip: the ParallelExpertEncoder online settings and the Sortformer streaming config
+        assert (pe.get("online_inference_length", 500), pe.get("chunk_left_context", 50), pe.get("chunk_right_context", 50),
+                pe.get("diar_spkcache_len", 188), pe.get("diar_fifo_len", 40), pe.get("diar_spkcache_update_period", 300),
+                pe.get("speaker_activity_threshold", 0.5)) == (500, 50, 50, 188, 40, 300, 0.5), \
+            "streaming settings differ from the constants in tools/mtmd/models/nemotron-transcribe.cpp"
+        sm = diar["sortformer_modules"]
+        assert (sm["spkcache_sil_frames_per_spk"], sm["pred_score_threshold"], sm["scores_boost_latest"], sm["strong_boost_rate"],
+                sm["weak_boost_rate"], sm["min_pos_scores_rate"], sm["use_learnable_sil_emb"]) == (1, 0.25, 0.05, 0.75, 1.5, 0.5, True), \
+            "speaker cache settings differ from the constants in tools/mtmd/models/nemotron-transcribe.cpp"
+
+        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.NEMOTRON_TRANSCRIBE)
+        self.gguf_writer.add_audio_num_mel_bins(asr["feat_in"])
+        self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)
+        self.gguf_writer.add_audio_subsampling_factor(asr["subsampling_factor"])
+        self.gguf_writer.add_audio_rope_dims(self._rope_dims(asr))
+        self.gguf_writer.add_audio_spk_block_count(spk["n_layers"])
+        self.gguf_writer.add_audio_spk_head_count(spk["n_heads"])
+
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        # ParallelExpertEncoder._build_sinusoid_position_encoding
+        pe = self.global_config["pe_encoder_config"]
+        n_spk = int(pe["diarization_model_cfg"]["sortformer_modules"]["num_spks"])
+        d = int(pe["asr_encoder_cfg"]["d_model"])
+        scale = float(pe.get("spk_kernel_scale", 1.0))
+        pos = torch.arange(n_spk, dtype=torch.float32)[:, None]
+        div = torch.exp(torch.arange(0, d, 2, dtype=torch.float32) * -(np.log(10000.0) / d))
+        kernel = torch.zeros(n_spk, d, dtype=torch.float32)
+        kernel[:, 0::2] = torch.sin(pos * div)
+        kernel[:, 1::2] = torch.cos(pos * div)
+        # stored as [d, n_spk]: the graph multiplies it with the speaker activity directly
+        yield self.format_tensor_name(gguf.MODEL_TENSOR.A_FUSION_SPK_KERNEL, suffix=""), (kernel * scale).T.contiguous()
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.A_SPK_HEAD_UPSAMPLE, None):
+            return gguf.GGMLQuantizationType.F16  # ggml_conv_1d kernel
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+        if not name.startswith("perception.") or "_extra_state" in name:
+            return None
+        # training-only head; the speaker model's unused front end
+        if name.startswith(("perception.encoder.diarization_model.sortformer_modules.hidden_to_spks.",
+                            "perception.encoder.diarization_model.preprocessor.")):
+            return None
+        return name, gen
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name == "perception.preprocessor.featurizer.fb":
+            data_torch = data_torch.squeeze(0)
+        yield from super().modify_tensors(data_torch, name, bid)

conversion/pockettts.py

@@ -5,6 +5,7 @@ from pathlib import Path
 from typing import Any, Iterable, TYPE_CHECKING
 
 import torch
+import torch.nn.functional as F
 
 if TYPE_CHECKING:
     from torch import Tensor
@@ -45,6 +46,102 @@ def _tensor_shapes(dir_model: Path) -> dict[str, tuple[int, ...]]:
         return {name: tuple(part[name].shape) for name in part.keys()}
 
 
+def _has_voice_encoder(dir_model: Path) -> bool:
+    # without voice cloning, the mimi encoder weights are all zero
+    part_names = ModelBase.get_model_part_names(dir_model, "model", ".safetensors")
+    with gguf.utility.SafetensorsLocal(dir_model / part_names[0]) as part:
+        nonzero = {bool(part[name].mmap_bytes().any()) for name in part.keys() if name.startswith("mimi.encoder")}
+    if len(nonzero) != 1:
+        raise ValueError(f"{dir_model}: the mimi encoder is partly zeroed out")
+    return nonzero.pop()
+
+
+def _load_safetensors(path: Path, prefixes: tuple[str, ...]) -> dict[str, Tensor]:
+    dtypes = {"F32": torch.float32, "BF16": torch.bfloat16, "I64": torch.int64}
+    with gguf.utility.SafetensorsLocal(path) as part:
+        return {name: torch.from_numpy(part[name].mmap_bytes()).view(dtypes[part[name].dtype]).reshape(part[name].shape)
+                for name in part.keys() if name.startswith(prefixes)}
+
+
+def _rope_pairs(x: Tensor, n_head: int, rope_theta: float) -> Tensor:
+    # ggml NORM mode rope, x: [n_pos, n_embd]
+    n_pos, n_embd = x.shape
+    head_dim = n_embd // n_head
+    angle = torch.arange(n_pos, dtype=x.dtype)[:, None, None] * \
+        torch.pow(rope_theta, -torch.arange(0, head_dim, 2, dtype=x.dtype) / head_dim)
+    x = x.view(n_pos, n_head, head_dim // 2, 2)
+    return torch.stack([x[..., 0] * angle.cos() - x[..., 1] * angle.sin(),
+                        x[..., 0] * angle.sin() + x[..., 1] * angle.cos()], -1).reshape(n_pos, n_embd)
+
+
+def _voice_prompt_embd(weights: dict[str, Tensor], kv: list[Tensor], n_head: int, rope_theta: float) -> tuple[Tensor, float]:
+    # the voice state only has K/V of each layer (kv[il] is [k/v, pos, n_embd]): layer 0 V gives each frame up to its scale, layer 1 V gives the scale
+    # returns the frames without bos_before_voice, and the max relative K/V error when the frames are prompted again
+    w = {name: t.double() for name, t in weights.items()}
+    kv = [t.double() for t in kv]
+    k0, v0, v1 = kv[0][0], kv[0][1], kv[1][1]
+    n_pos, n_embd = v0.shape
+    head_dim = n_embd // n_head
+
+    def p(il: int, name: str) -> Tensor:
+        return w[f"flow_lm.transformer.layers.{il}.{name}"]
+
+    def norm(x: Tensor, il: int, name: str) -> Tensor:
+        return F.layer_norm(x, (n_embd,), p(il, f"{name}.weight"), p(il, f"{name}.bias"), 1e-5)
+
+    def heads(t: Tensor) -> Tensor:
+        return t.reshape(n_pos, n_head, head_dim).transpose(0, 1)
+
+    def rows(t: Tensor) -> Tensor:
+        return t.transpose(0, 1).reshape(n_pos, n_embd)
+
+    def attn(il: int, h: Tensor, k: Tensor, v: Tensor) -> Tensor:
+        q = _rope_pairs(h @ p(il, "self_attn.in_proj.weight").chunk(3)[0].T, n_head, rope_theta)
+        return rows(F.scaled_dot_product_attention(heads(q), heads(k), heads(v), is_causal=True)) @ p(il, "self_attn.out_proj.weight").T
+
+    def ffn(il: int, x: Tensor) -> Tensor:
+        return F.gelu(norm(x, il, "norm2") @ p(il, "linear1.weight").T) @ p(il, "linear2.weight").T
+
+    wv = p(0, "self_attn.in_proj.weight").chunk(3)[2]
+    h = torch.linalg.solve(wv, v0.T).T                   # norm1(x)
+    z = (h - p(0, "norm1.bias")) / p(0, "norm1.weight")  # (x - mean(x)) / std(x)
+
+    # fit z on [speaker_proj, 1] to get the latent direction back
+    proj = w["flow_lm.speaker_proj_weight"]
+    fit = torch.linalg.lstsq(torch.cat([proj, torch.ones(n_embd, 1, dtype=proj.dtype)], 1), z[1:].T).solution
+    frames = (proj @ fit[:-1]).T
+
+    a0 = attn(0, h, k0, v0)[1:]
+    wv1 = p(1, "self_attn.in_proj.weight").chunk(3)[2]
+
+    def err(scale: Tensor) -> Tensor:  # per frame relative error of layer 1 V
+        x = scale[:, None] * frames + a0
+        x = x + ffn(0, x)
+        return (norm(x, 1, "norm1") @ wv1.T - v1[1:]).norm(dim=1) / v1[1:].norm(dim=1)
+
+    grid = torch.logspace(-3, 3, 49, dtype=torch.float64)
+    best = torch.stack([err(s.expand(n_pos - 1)) for s in grid]).argmin(0)
+    lo, hi = grid[(best - 1).clamp(min=0)], grid[(best + 1).clamp(max=len(grid) - 1)]
+    for _ in range(40):  # golden section search, all frames at once
+        m1, m2 = hi - 0.618 * (hi - lo), lo + 0.618 * (hi - lo)
+        left = err(m1) < err(m2)
+        lo, hi = torch.where(left, lo, m1), torch.where(left, m2, hi)
+    frames = (lo + hi)[:, None] / 2 * frames
+
+    # prompt the frames again and compare K/V with the shipped state
+    x = torch.cat([w["flow_lm.bos_before_voice"].reshape(1, n_embd), frames])
+    drift = 0.0
+    for il, (k_ref, v_ref) in enumerate(kv):
+        h = norm(x, il, "norm1")
+        _, wk, wv = p(il, "self_attn.in_proj.weight").chunk(3)
+        k, v = _rope_pairs(h @ wk.T, n_head, rope_theta), h @ wv.T
+        for a, b in ((k, k_ref), (v, v_ref)):
+            drift = max(drift, ((a - b).norm(dim=1) / b.norm(dim=1)).median().item())
+        x = x + attn(il, h, k, v)
+        x = x + ffn(il, x)
+    return frames, drift
+
+
 @ModelBase.register_hparams_loader(lambda dir_model: "flow_lm.bos_emb" in _tensor_shapes(dir_model))
 def _load_hparams(dir_model: Path) -> dict[str, Any]:
     logger.info("gguf: detected pocket-tts checkpoint, deriving hparams from tensor shapes")
@@ -194,22 +291,27 @@ class PocketTTSMmprojModel(MmprojModel):
         (gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_Q, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_K, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_V),
     )
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.has_audio_encoder = _has_voice_encoder(self.dir_model)
+
     def set_gguf_parameters(self):
         self.gguf_writer.add_file_type(self.ftype)
         assert self.hparams_audio is not None
 
         # voice-prompt encoder: mimi encoder + speaker_proj
-        self.gguf_writer.add_clip_has_audio_encoder(True)
-        # note: the 24kHz sample rate is hardcoded on the clip.cpp side, like the other audio models
-        self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_SPKENC)
-        self.gguf_writer.add_audio_projection_dim(self.n_embd_text)
-        self.gguf_writer.add_audio_block_count(self.hparams_audio["num_hidden_layers"])
-        self.gguf_writer.add_audio_embedding_length(self.hparams_audio["hidden_size"])
-        self.gguf_writer.add_audio_feed_forward_length(self.hparams_audio["intermediate_size"])
-        self.gguf_writer.add_audio_head_count(self.hparams_audio["num_attention_heads"])
-        self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)
-        # mimi convolves the waveform directly, it is passed around as a 1-row "mel"
-        self.gguf_writer.add_audio_num_mel_bins(1)
+        if self.has_audio_encoder:
+            self.gguf_writer.add_clip_has_audio_encoder(True)
+            # note: the 24kHz sample rate is hardcoded on the clip.cpp side, like the other audio models
+            self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_SPKENC)
+            self.gguf_writer.add_audio_projection_dim(self.n_embd_text)
+            self.gguf_writer.add_audio_block_count(self.hparams_audio["num_hidden_layers"])
+            self.gguf_writer.add_audio_embedding_length(self.hparams_audio["hidden_size"])
+            self.gguf_writer.add_audio_feed_forward_length(self.hparams_audio["intermediate_size"])
+            self.gguf_writer.add_audio_head_count(self.hparams_audio["num_attention_heads"])
+            self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)
+            # mimi convolves the waveform directly, it is passed around as a 1-row "mel"
+            self.gguf_writer.add_audio_num_mel_bins(1)
 
         # generation: flow-matching decoder + mimi decoder
         # the SEANet and flow net hparams are constant across the family, clip.cpp holds them
@@ -229,12 +331,57 @@ class PocketTTSMmprojModel(MmprojModel):
         # conv1d/conv1d_dw kernels must be F16, ggml_conv_1d(_dw) has no BF16 path
         if ".seanet." in new_name or new_name in ("a.downsample.conv.weight", "a.gen.wav.upsample.weight"):
             return gguf.GGMLQuantizationType.F16
+        # voices are read on the host
+        if new_name.startswith("voice."):
+            return gguf.GGMLQuantizationType.F32
         return False
 
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        # embeddings/<name>.safetensors hold the K/V state after the voice prompt (see get_state_for_audio_prompt() in the reference), turn it back into prompt embeddings
+        paths = sorted((self.dir_model / "embeddings").glob("*.safetensors"))
+        if not paths:
+            return
+        weights = _load_safetensors(self.dir_model / "model.safetensors",
+                                    ("flow_lm.transformer.layers.", "flow_lm.speaker_proj_weight", "flow_lm.bos_before_voice"))
+        n_layer = self.global_config["num_hidden_layers"]
+        for path in paths:
+            state = _load_safetensors(path, ("transformer.layers.",))
+            kv = []
+            for il in range(n_layer):
+                prefix = f"transformer.layers.{il}.self_attn/"
+                if any(prefix + k not in state for k in ("cache", "offset")):
+                    break
+                n_pos = int(state[prefix + "offset"].item())
+                # older packs have no pad entry
+                if prefix + "pad" in state and int(state[prefix + "pad"].item()) != 0:
+                    break
+                kv.append(state[prefix + "cache"][:, 0, :n_pos].flatten(2))  # [k/v, pos, n_embd]
+            if len(kv) != n_layer or any(t.shape != kv[0].shape for t in kv) or kv[0].shape[1] < 2 or \
+                    kv[0].shape[2] != self.global_config["hidden_size"]:
+                logger.warning(f"voice {path.stem}: skipped, not an unpadded voice state of this backbone")
+                continue
+            embd, drift = _voice_prompt_embd(weights, kv, self.global_config["num_attention_heads"],
+                                             self.global_config["rope_theta"])
+            # some packs ship a state made with other weights, the reference still uses it: keep up to 50% error, warn above 5%
+            if drift > 0.5:
+                logger.warning(f"voice {path.stem}: skipped, its state does not belong to these weights (error {drift:.1e})")
+                continue
+            log = logger.warning if drift > 0.05 else logger.info
+            log(f"voice {path.stem}: {embd.shape[0]} frames, prompting them again gives the shipped state within {drift:.1e}")
+            yield (f"voice.{path.stem}", embd.float())
+
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
         del bid  # the block index of the mimi transformers is parsed here, not by the base class
         T = gguf.MODEL_TENSOR
 
+        if name.startswith("voice."):
+            yield (name, data_torch)
+            return
+
+        if not self.has_audio_encoder and (name.startswith(("mimi.encoder", "mimi.downsample."))
+                                           or name == "flow_lm.speaker_proj_weight"):
+            return  # voice-prompt encoder
+
         if name in ("flow_lm.bos_emb", "flow_lm.bos_before_voice", "flow_lm.conditioner.embed.weight"):
             return  # folded into the backbone embedding table
         if name.startswith("flow_lm.transformer.") or name.startswith("flow_lm.out_norm."):

conversion/qwen.py

@@ -83,6 +83,15 @@ class Qwen2Model(TextModel):
 class Qwen2MoeModel(TextModel):
     model_arch = gguf.MODEL_ARCH.QWEN2MOE
 
+    def _map_fp8_weight_names(self, name: str) -> tuple[str, ...]:
+        if name.removesuffix(".weight").endswith(".mlp.experts.gate_up_proj"):
+            bid = next(int(part) for part in self.map_tensor_name(name).split(".") if part.isdecimal())
+            return (
+                self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_EXP, bid),
+                self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP_EXP, bid),
+            )
+        return super()._map_fp8_weight_names(name)
+
     def set_gguf_parameters(self):
         super().set_gguf_parameters()
         if (moe_intermediate_size := self.hparams.get("moe_intermediate_size")) is not None:
@@ -382,6 +391,16 @@ class _QwenMtpMixin:
 class Qwen3NextModel(_QwenMtpMixin, Qwen2MoeModel):
     model_arch = gguf.MODEL_ARCH.QWEN3NEXT
 
+    def _map_fp8_weight_names(self, name: str) -> tuple[str, ...]:
+        if name.endswith(".linear_attn.in_proj_qkvz.weight"):
+            bid = next(int(part) for part in self.map_tensor_name(name).split(".") if part.isdecimal())
+            # Both split projections keep the original per-tensor scales.
+            return (
+                self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_QKV, bid),
+                self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_GATE, bid),
+            )
+        return super()._map_fp8_weight_names(name)
+
     def set_gguf_parameters(self):
         super().set_gguf_parameters()
         self.gguf_writer.add_ssm_conv_kernel(self.hparams["linear_conv_kernel_dim"])

convert_hf_to_gguf.py

@@ -58,7 +58,7 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument(
         "--outtype", type=str, choices=["f32", "f16", "bf16", "q8_0", "tq1_0", "tq2_0", "auto"], default="auto",
-        help="output format - use f32 for float32, f16 for float16, bf16 for bfloat16, q8_0 for Q8_0, tq1_0 or tq2_0 for ternary, and auto for the highest-fidelity 16-bit float type",
+        help="output format - use f32 for float32, f16 for float16, bf16 for bfloat16, q8_0 for Q8_0, tq1_0 or tq2_0 for ternary, and auto to preserve supported FP8 weights and otherwise select the highest-fidelity 16-bit float type",
     )
     parser.add_argument(
         "--bigendian", action="store_true",
@@ -127,7 +127,7 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument(
         "--dspark", action="store_true",
-        help="Export only the DeepSeek-V4 DSpark draft tensors as a separate GGUF.",
+        help="Export only the DeepSeek-V4 or Mistral Large 4 DSpark draft tensors as a separate GGUF.",
     )
     parser.add_argument(
         "--mistral-format", action="store_true",
@@ -255,6 +255,9 @@ def main() -> None:
             assert hparams.get("vision_encoder") is not None, "This model does not support multimodal"
             from conversion.pixtral import PixtralModel
             model_class = PixtralModel
+        elif hparams.get("moe") is not None and hparams.get("dsa") is not None:
+            from conversion.mistral4 import MistralLarge4Model
+            model_class = MistralLarge4Model
         elif hparams.get("moe") is not None:
             from conversion.mistral import MistralMoeModel
             model_class = MistralMoeModel
@@ -267,10 +270,16 @@ def main() -> None:
             sys.exit(1)
 
         if args.dspark:
-            if is_mistral_format or model_architecture not in ("DeepseekV4ForCausalLM", "DeepseekV41ForCausalLM"):
-                logger.error("--dspark is only supported for DeepseekV4ForCausalLM and DeepseekV41ForCausalLM")
+            if is_mistral_format or model_architecture == "MistralLarge4ForConditionalGeneration":
+                if model_class.__name__ != "MistralLarge4Model":
+                    logger.error("--dspark is only supported for Mistral Large 4 among Mistral-format models")
+                    sys.exit(1)
+                from conversion.mistral4 import MistralLarge4DSparkModel
+                model_class = MistralLarge4DSparkModel
+            elif model_architecture not in ("DeepseekV4ForCausalLM", "DeepseekV41ForCausalLM"):
+                logger.error("--dspark is only supported for DeepseekV4ForCausalLM, DeepseekV41ForCausalLM and Mistral Large 4")
                 sys.exit(1)
-            if model_architecture == "DeepseekV41ForCausalLM":
+            elif model_architecture == "DeepseekV41ForCausalLM":
                 from conversion.deepseek41 import DeepseekV41DSparkModel
                 model_class = DeepseekV41DSparkModel
             else:

docs/backend/CANN.md

@@ -161,7 +161,7 @@ lspci -n | grep -Eo '19e5:d[0-9a-f]{3}' | cut -d: -f2
 | FP16                   | Support | Support |
 | Q8_0                   | Support | Partial |
 | Q4_0                   | Support | Partial |
-| BF16                   | Support |         |
+| BF16                   | Support |    No   |
 
 > **310P note**
 > - `Q8_0`: data transform / buffer path is implemented, and `GET_ROWS` is supported, but quantized `MUL_MAT` / `MUL_MAT_ID` are not supported.
@@ -274,7 +274,9 @@ cmake --build build --config release
 
     **Notes**:
 
-      - CANN backend only supports FP16/Q4_0/Q8_0 models currently.
+      - On 910B the CANN backend supports FP16, BF16 and Q4_0/Q8_0 models. On 310P, BF16 is not
+        supported and quantized `MUL_MAT` / `MUL_MAT_ID` are not implemented, so only FP16 models
+        are supported (see [DataType Supports](#datatype-supports)).
 
 2. **Launch inference**
 
@@ -355,3 +357,7 @@ Enable ACL graph execution during the prefill stage, default is false. This opti
 ### GGML_CANN_OPERATOR_FUSION
 
 Enable operator fusion during computation, default is false. This option fuses compatible operators (e.g., ADD + RMS_NORM) to reduce overhead and improve performance.
+
+### GGML_CANN_NO_PINNED
+
+Disable pinned (host-locked) host buffers. When it is set (any value), `aclrtMallocHost` is not called: the CANN host buffer type allocates regular CPU memory instead and the backend reports no `host_buffer` capability, so host-side staging uses pageable memory. Use it when pinned host allocation fails, or when the pinned memory is needed by other processes on the host.

docs/build-s390x.md

@@ -249,7 +249,7 @@ IBM VXE/VXE2 SIMD acceleration depends on the BLAS implementation. It is strongl
 | MXFP4      | ✅           | ❓    | ❓     |
 | Q5_0       | ✅           | ❓    | ❓     |
 | Q5_1       | ✅           | ❓    | ❓     |
-| Q8_0       | ✅           | ❓    | ❓     |
+| Q8_0       | ✅           | ✅    | ❓     |
 | Q2_K       | 🚫           | ❓    | ❓     |
 | Q3_K       | ✅           | ❓    | ❓     |
 | Q4_K       | ✅           | ❓    | ❓     |
@@ -266,11 +266,11 @@ IBM VXE/VXE2 SIMD acceleration depends on the BLAS implementation. It is strongl
 | IQ1_M      | 🚫           | ❓    | ❓     |
 | IQ4_NL     | ✅           | ❓    | ❓     |
 | IQ4_XS     | ✅           | ❓    | ❓     |
-| FP32->FP16 | 🚫           | ❓    | ❓     |
+| FP32->FP16 | ✅           | ❓    | ❓     |
 | FP16->FP32 | 🚫           | ❓    | ❓     |
 
 -   ✅ - acceleration available
 -   🚫 - acceleration unavailable, will still run using scalar implementation
 -   ❓ - acceleration unknown, please contribute if you can test it yourself
 
-Last Updated by **Aaron Teo (aaron.teo1@ibm.com)** on Sep 8, 2026.
+Last Updated by **Aaron Teo (aaron.teo1@ibm.com)** on Nov 3, 2026.

docs/build.md

@@ -481,7 +481,7 @@ cmake --build build --config Release
 Now you can load the model in conversation mode using `Vulkan`
 
 ```sh
-build/bin/Release/llama-cli -m "[PATH TO MODEL]" -ngl 100 -c 16384 -t 10 -n -2 -cnv
+build/bin/Release/llama-cli -m "[PATH TO MODEL]" -ngl 100 -c 16384 -t 10 -n -2
 ```
 
 **MSYS2**

docs/ops.md

@@ -14,22 +14,22 @@ Legend:
 
 | Operation | BLAS | CANN | CPU | CUDA | ET | HTP | MTL | OpenCL | SYCL | Vulkan | WebGPU | ZenDNN | zDNN |
 |-----------|------|------|------|------|------|------|------|------|------|------|------|------|------|
-|                              ABS | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              ABS | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                              ACC | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                              ADD | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              ADD | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             ADD1 | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                           ADD_ID | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                           ARANGE | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                           ARGMAX | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                           ARANGE | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                           ARGMAX | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          ARGSORT | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             CEIL | ❌ | ❌ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             CEIL | ❌ | ❌ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                            CLAMP | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        COL2IM_1D | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                           CONCAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                             CONT | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
-|                       CONV_1D_DW | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ |
+|                       CONV_1D_DW | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ |
 |                          CONV_2D | ❌ | ❌ | 🟡 | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ |
-|                       CONV_2D_DW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                       CONV_2D_DW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          CONV_3D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                CONV_TRANSPOSE_1D | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                CONV_TRANSPOSE_2D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -38,49 +38,49 @@ Legend:
 |                              CPY | ❌ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
 |               CROSS_ENTROPY_LOSS | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |          CROSS_ENTROPY_LOSS_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                           CUMSUM | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             DIAG | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                           CUMSUM | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             DIAG | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                    DIAG_MASK_INF | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                              DIV | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              DIV | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                     DSV4_HC_COMB | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                     DSV4_HC_POST | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                      DSV4_HC_PRE | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                              DUP | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                              ELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                              EXP | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                            EXPM1 | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             FILL | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              ELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              EXP | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                            EXPM1 | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             FILL | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                   FLASH_ATTN_EXT | ❌ | 🟡 | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
-|                            FLOOR | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                  GATED_DELTA_NET | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | 🟡 | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                            FLOOR | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                  GATED_DELTA_NET | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                GATED_LINEAR_ATTN | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                            GEGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                        GEGLU_ERF | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                      GEGLU_QUICK | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                            GEGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                        GEGLU_ERF | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                      GEGLU_QUICK | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             GELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                         GELU_ERF | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                       GELU_QUICK | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                         GET_ROWS | ❌ | 🟡 | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ❌ | ❌ |
 |                    GET_ROWS_BACK | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ |
 |                       GROUP_NORM | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                      HARDSIGMOID | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                        HARDSWISH | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                      HARDSIGMOID | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                        HARDSWISH | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                           IM2COL | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                        IM2COL_3D | ❌ | ❌ | ✅ | ✅ | ❌ | 🟡 | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                          L2_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | 🟡 | ❌ | ❌ |
+|                        IM2COL_3D | ❌ | ❌ | ✅ | ✅ | ❌ | 🟡 | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                          L2_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                       LEAKY_RELU | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | 🟡 | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                LIGHTNING_INDEXER | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                LIGHTNING_INDEXER | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | 🟡 | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                              LOG | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             MEAN | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ | ❌ |
-|                              MUL | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              MUL | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          MUL_MAT | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 |
-|                 MUL_MAT_HADAMARD | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 |
+|                 MUL_MAT_HADAMARD | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | 🟡 |
 |                       MUL_MAT_ID | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | 🟡 | ❌ |
-|                  MUL_MAT_ID_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                  MUL_MAT_ID_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                     MUL_MAT_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                     MUL_MAT_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                              NEG | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                  MUL_MAT_ID_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | 🟡 | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                  MUL_MAT_ID_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | 🟡 | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                     MUL_MAT_W4A4 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | 🟡 | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                     MUL_MAT_W4A8 | ❌ | ❌ | ✅ | ❌ | ❌ | 🟡 | ❌ | 🟡 | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                              NEG | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                   OPT_STEP_ADAMW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                     OPT_STEP_SGD | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -89,45 +89,46 @@ Legend:
 |                   PAD_REFLECT_1D | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                          POOL_1D | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                          POOL_2D | ❌ | 🟡 | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                            REGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                            REGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             RELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                           REPEAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                      REPEAT_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                         RMS_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                         RMS_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                    RMS_NORM_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                             ROLL | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                             ROPE | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        ROPE_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                            ROUND | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                            ROUND | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        RWKV_WKV6 | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                        RWKV_WKV7 | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                            SCALE | ❌ | 🟡 | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                              SET | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ |
+|                            SCALE | ❌ | 🟡 | ✅ | ✅ | ✅ | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              SET | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ✅ | ❌ | ❌ |
 |                         SET_ROWS | ❌ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
-|                              SGN | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              SGN | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          SIGMOID | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             SILU | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        SILU_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
 |                              SIN | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                         SOFTPLUS | ❌ | ❌ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                         SOFTPLUS | ❌ | ❌ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                         SOFT_MAX | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                    SOFT_MAX_BACK | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ | ❌ | 🟡 | ✅ | ❌ | ❌ | ❌ |
-|                        SOLVE_TRI | ❌ | ❌ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ |
+|                        SOLVE_TRI | ❌ | ❌ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ❌ | ❌ |
 |                              SQR | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             SQRT | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                         SSM_CONV | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                         SSM_SCAN | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | 🟡 | 🟡 | ✅ | ❌ | ❌ |
-|                             STEP | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                              SUB | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                         SSM_SCAN | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ❌ | ❌ |
+|                             STEP | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                              SUB | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                              SUM | ❌ | 🟡 | ✅ | 🟡 | ❌ | 🟡 | 🟡 | ❌ | 🟡 | 🟡 | 🟡 | ❌ | ❌ |
 |                         SUM_ROWS | ❌ | ✅ | ✅ | 🟡 | ❌ | 🟡 | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ |
-|                           SWIGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                     SWIGLU_CLAMP | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
+|                           SWIGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                     SWIGLU_CLAMP | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                       SWIGLU_OAI | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             TANH | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                             TANH | ❌ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |               TIMESTEP_EMBEDDING | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                            TOP_K | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | 🟡 | ✅ | ❌ | ❌ |
-|                              TRI | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                            TRUNC | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                            TOP_K | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ |
+|                       TOP_K_MASK | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                              TRI | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                            TRUNC | ❌ | ❌ | ✅ | 🟡 | 🟡 | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          UPSCALE | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                            XIELU | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |

docs/ops/WebGPU.csv

@@ -12326,8 +12326,16 @@
 "WebGPU: WebGPU","PAD_REFLECT_1D","type=f32,ne_a=[512,34,2,1],pad_0=10,pad_1=9","support","0","no","WebGPU"
 "WebGPU: WebGPU","PAD_REFLECT_1D","type=f32,ne_a=[3000,384,4,1],pad_0=10,pad_1=9","support","0","no","WebGPU"
 "WebGPU: WebGPU","ROLL","shift0=3,shift1=-2,shift3=1,shift4=-1","support","0","no","WebGPU"
-"WebGPU: WebGPU","ARANGE","type=f32,start=0.000000,stop=10.000000,step=1.000000","support","0","no","WebGPU"
-"WebGPU: WebGPU","ARANGE","type=f32,start=0.000000,stop=1048576.000000,step=1.000000","support","0","no","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=0.000000,stop=10.000000,step=1.000000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=3.000000,stop=3.500000,step=1.000000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=-10.000000,stop=10.000000,step=0.300000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=-1.000000,stop=62.750000,step=0.250000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=-1.000000,stop=63.000000,step=0.250000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=-1.000000,stop=63.250000,step=0.250000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=-1.000000,stop=254.750000,step=0.250000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=-1.000000,stop=255.000000,step=0.250000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=-1.000000,stop=255.250000,step=0.250000","support","1","yes","WebGPU"
+"WebGPU: WebGPU","ARANGE","type=f32,start=0.000000,stop=1048576.000000,step=1.000000","support","1","yes","WebGPU"
 "WebGPU: WebGPU","TIMESTEP_EMBEDDING","type=f32,ne_a=[2,1,1,1],dim=320,max_period=10000","support","0","no","WebGPU"
 "WebGPU: WebGPU","LEAKY_RELU","type=f32,ne_a=[10,5,4,3],negative_slope=0.100000","support","0","no","WebGPU"
 "WebGPU: WebGPU","CUMSUM","type=f32,ne=[10,5,4,3]","support","1","yes","WebGPU"

ggml/include/ggml-backend.h

@@ -70,7 +70,8 @@ extern "C" {
     GGML_API void                           ggml_backend_buffer_reset         (ggml_backend_buffer_t buffer);
 
     // tensor copy between different backends
-    GGML_API void ggml_backend_tensor_copy(const struct ggml_tensor * src, struct ggml_tensor * dst);
+    GGML_API void ggml_backend_tensor_copy  (const struct ggml_tensor * src, struct ggml_tensor * dst);
+    GGML_API void ggml_backend_tensor_copy_n(const struct ggml_tensor * const * srcs, struct ggml_tensor * const * dsts, size_t n_tensors);
 
     //
     // Backend (stream)
@@ -95,6 +96,8 @@ extern "C" {
     GGML_API void ggml_backend_tensor_get   (const struct ggml_tensor * tensor,       void * data, size_t offset, size_t size);
     GGML_API void ggml_backend_tensor_set_2d(      struct ggml_tensor * tensor, const void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data);
     GGML_API void ggml_backend_tensor_get_2d(const struct ggml_tensor * tensor,       void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data);
+    GGML_API void ggml_backend_tensor_set_n (      struct ggml_tensor * const * tensors, const void * data, const size_t * offsets, const size_t * sizes, const size_t * data_offsets, size_t n_tensors);
+    GGML_API void ggml_backend_tensor_get_n (const struct ggml_tensor * const * tensors,       void * data, const size_t * offsets, const size_t * sizes, const size_t * data_offsets, size_t n_tensors);
     GGML_API void ggml_backend_tensor_memset(      struct ggml_tensor * tensor,     uint8_t value, size_t offset, size_t size);
 
     GGML_API void ggml_backend_synchronize(ggml_backend_t backend);
@@ -208,6 +211,18 @@ extern "C" {
 
     // Context management and operations for faster communication between backends, used for tensor parallelism (meta backend)
     typedef void * (*ggml_backend_comm_init_t)(ggml_backend_t * backends, size_t n_backends);
+
+    // Builds the communicator for the case where each rank is a separate process, so a process
+    // owns one backend of the communicator rather than all of them.
+    //
+    // The caller provides bootstrap_callback, which moves bytes between the ranks before the
+    // communicator exists. Each rank passes size bytes in send_buf, and recv_buf receives
+    // world_size * size bytes in rank order. The backend decides what to exchange and how many
+    // times to call it, and every rank must make the same calls in the same order.
+    // comm_init_rank blocks until all world_size ranks have joined.
+    typedef bool   (*ggml_backend_comm_bootstrap_callback)(void * user_data, const void * send_buf, void * recv_buf, size_t size);
+    typedef void * (*ggml_backend_comm_init_rank_t)(ggml_backend_t backend, int rank, int world_size, ggml_backend_comm_bootstrap_callback bootstrap_callback, void * bootstrap_callback_data);
+
     typedef void   (*ggml_backend_comm_free_t)(void * comm_ctx);
     typedef bool   (*ggml_backend_comm_allreduce_tensor_t)(void * comm_ctx, struct ggml_tensor ** tensors);
 

ggml/include/ggml.h

@@ -381,6 +381,8 @@ extern "C" {
     GGML_API void        ggml_fp32_to_bf16_row_ref(const float *, ggml_bf16_t *, int64_t);
     GGML_API void        ggml_fp32_to_bf16_row(const float *, ggml_bf16_t *, int64_t);
 
+    typedef struct { uint8_t bits; } ggml_fp8_e4m3_t;
+
     struct ggml_object;
     struct ggml_context;
     struct ggml_cgraph;
@@ -430,7 +432,8 @@ extern "C" {
         GGML_TYPE_NVFP4   = 40, // NVFP4 (4 blocks, E4M3 scale)
         GGML_TYPE_Q1_0    = 41,
         GGML_TYPE_Q2_0    = 42,
-        GGML_TYPE_COUNT   = 43,
+        GGML_TYPE_F8_E4M3 = 43,
+        GGML_TYPE_COUNT   = 44,
     };
 
     // [TAG_GGML_PREC]
@@ -486,6 +489,7 @@ extern "C" {
         GGML_FTYPE_MOSTLY_NVFP4   = 26, // except 1d tensors
         GGML_FTYPE_MOSTLY_Q1_0    = 27, // except 1d tensors
         GGML_FTYPE_MOSTLY_Q2_0    = 28, // except 1d tensors
+        GGML_FTYPE_MOSTLY_F8_E4M3 = 29, // except 1d tensors
     };
 
     // available tensor operations:
@@ -2512,6 +2516,24 @@ extern "C" {
             float                 max_bias,
             float                 logit_softcap);
 
+    // same as ggml_flash_attn_ext, but each q slice attends to its own list of K/V rows
+    // q:       [n_embd_k, n_batch, n_head,    ne3]
+    // k:       [n_embd_k, n_rows,  n_head_kv, 1  ] shared by all q slices
+    // v:       [n_embd_v, n_rows,  n_head_kv, 1  ]
+    // mask:    [n_kv,     n_batch, ne32,      ne3]
+    // kv_rows: [n_kv,     ne3] I32, mask column i of slice i3 uses row kv_rows[i, i3] of k and v
+    // negative rows are padding, their mask entries must be -inf
+    GGML_API struct ggml_tensor * ggml_flash_attn_ext_rows(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * q,
+            struct ggml_tensor  * k,
+            struct ggml_tensor  * v,
+            struct ggml_tensor  * mask,
+            struct ggml_tensor  * kv_rows,
+            float                 scale,
+            float                 max_bias,
+            float                 logit_softcap);
+
     GGML_DEPRECATED(GGML_API void ggml_flash_attn_ext_set_prec(
             struct ggml_tensor * a,
             enum ggml_prec       prec),

ggml/src/ggml-backend-impl.h

@@ -47,6 +47,9 @@ extern "C" {
     // Backend buffer
     //
 
+    // TODO: remove `set_tensor`, `get_tensor`, `set_tensor_2d`, `get_tensor_2d` and `cpy_tensor` from the buffer interface
+    //       they are superseded by the `set_tensor_n`, `get_tensor_n`, `cpy_tensor_n` interface
+    //       ref: https://github.com/ggml-org/llama.cpp/pull/30881#discussion_r4400299490
     struct ggml_backend_buffer_i {
         // (optional) free the buffer
         void         (*free_buffer)  (ggml_backend_buffer_t buffer);
@@ -68,6 +71,11 @@ extern "C" {
         void         (*clear)        (ggml_backend_buffer_t buffer, uint8_t value);
         // (optional) reset any internal state due to tensor initialization, such as tensor extras
         void         (*reset)        (ggml_backend_buffer_t buffer);
+        // (optional) batched tensor data copies
+        void         (*set_tensor_n) (ggml_backend_buffer_t buffer,       struct ggml_tensor * const * tensors, const void * const * data, const size_t * offsets, const size_t * sizes, size_t n_tensors);
+        void         (*get_tensor_n) (ggml_backend_buffer_t buffer, const struct ggml_tensor * const * tensors,       void * const * data, const size_t * offsets, const size_t * sizes, size_t n_tensors);
+        // (optional) batched tensor copies: returns true if all copies were performed, false if unsupported and no copies were performed
+        bool         (*cpy_tensor_n) (ggml_backend_buffer_t buffer, const struct ggml_tensor * const * srcs, struct ggml_tensor * const * dsts, size_t n_tensors);
     };
 
     struct ggml_backend_buffer {

ggml/src/ggml-backend-meta.cpp

@@ -489,6 +489,32 @@ static struct ggml_tensor * ggml_backend_meta_buffer_simple_tensor(const struct
     return it->second[index];
 }
 
+// bytes spanned by all dims except skip_dim
+static size_t ggml_backend_meta_inner_extent(const struct ggml_tensor * tensor, int skip_dim) {
+    size_t extent = 0;
+    for (int i = 0; i < GGML_MAX_DIMS; i++) {
+        if (i != skip_dim && tensor->ne[i] > 1) {
+            extent += (size_t) (tensor->ne[i] - 1) * tensor->nb[i];
+        }
+    }
+    return extent;
+}
+
+// the view dim that holds the split axis of src, or -1
+// all other dims of src and view must stay inside one slab of that axis
+static int ggml_backend_meta_view_slab_dim(const struct ggml_tensor * view, const struct ggml_tensor * src, int axis) {
+    const size_t slab = src->nb[axis];
+    if (ggml_backend_meta_inner_extent(src, axis) >= slab) {
+        return -1;
+    }
+    for (int dim = 0; dim < GGML_MAX_DIMS; dim++) {
+        if (view->ne[dim] == src->ne[axis] && view->nb[dim] == slab) {
+            return ggml_backend_meta_inner_extent(view, dim) < slab ? dim : -1;
+        }
+    }
+    return -1;
+}
+
 static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(const struct ggml_tensor * tensor, bool assume_sync);
 
 static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
@@ -700,9 +726,39 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
             }
             GGML_ABORT("fatal error");
         }
+        // permuted view: split the view dim that spans the src split axis, other dims must be inside one element of it or outside the span
+        if (ggml_is_permuted(tensor) && !ggml_is_permuted(tensor->src[0]) && axis >= 0 && axis < GGML_MAX_DIMS-1) {
+            const size_t nb_axis = tensor->src[0]->nb[axis];
+            const size_t nb_next = tensor->src[0]->nb[axis+1];
+            for (int dim = 0; dim < GGML_MAX_DIMS; dim++) {
+                if (tensor->ne[dim] == 1 || tensor->nb[dim] < nb_axis || tensor->nb[dim]*tensor->ne[dim] != nb_next) {
+                    continue;
+                }
+                bool others_ok = true;
+                for (int i = 0; i < GGML_MAX_DIMS; i++) {
+                    if (i == dim || tensor->ne[i] == 1) {
+                        continue;
+                    }
+                    others_ok = others_ok && (tensor->nb[i]*tensor->ne[i] <= tensor->nb[dim] || tensor->nb[i] % nb_next == 0);
+                }
+                if (others_ok) {
+                    return {ggml_backend_meta_split_axis(dim), {0}, {1}, 1};
+                }
+            }
+        }
         if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED || src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) {
             return src_ss[0];
         }
+        // a view that only regroups dims inside one slab of the split axis keeps the split,
+        // the per-device strides are only rescaled above nb[axis] (e.g. the KV streams of a permuted q)
+        if (axis >= 0 && axis < GGML_MAX_DIMS && tensor->view_offs == tensor->src[0]->view_offs) {
+            const int dim = ggml_backend_meta_view_slab_dim(tensor, tensor->src[0], axis);
+            if (dim >= 0) {
+                ggml_backend_meta_split_state ret = src_ss[0];
+                ret.axis = ggml_backend_meta_split_axis(dim);
+                return ret;
+            }
+        }
         GGML_ABORT("view of permuted tensor not implemented");
         //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1};
     };
@@ -1000,8 +1056,13 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
                 split_state = handle_pad(src_ss);
             } break;
             case GGML_OP_PAD_REFLECT_1D:
-            case GGML_OP_ROLL:
-            case GGML_OP_ARANGE:
+            case GGML_OP_ROLL: {
+                split_state = handle_generic(src_ss, /*scalar_only =*/ true);
+            } break;
+            case GGML_OP_ARANGE: {
+                // no sources, every device computes the full range
+                split_state = {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1};
+            } break;
             case GGML_OP_TIMESTEP_EMBEDDING: {
                 split_state = handle_generic(src_ss, /*scalar_only =*/ true);
             } break;
@@ -1699,6 +1760,9 @@ static const ggml_backend_buffer_i ggml_backend_meta_buffer_iface = {
     /* .cpy_tensor      = */ nullptr,
     /* .clear           = */ ggml_backend_meta_buffer_clear,
     /* .reset           = */ ggml_backend_meta_buffer_reset,
+    /* .set_tensor_n    = */ NULL,
+    /* .get_tensor_n    = */ NULL,
+    /* .cpy_tensor_n    = */ NULL,
 };
 
 bool ggml_backend_buffer_is_meta(ggml_backend_buffer_t buf) {

ggml/src/ggml-backend.cpp

@@ -402,6 +402,62 @@ size_t ggml_backend_get_max_size(ggml_backend_t backend) {
     return ggml_backend_buft_get_max_size(ggml_backend_get_default_buffer_type(backend));
 }
 
+template <typename cb_get_buffer>
+static void ggml_backend_tensor_buft_groups(size_t n_tensors, cb_get_buffer && get_buffer, std::vector<std::vector<size_t>> & groups) {
+    std::unordered_map<ggml_backend_buffer_type_t, size_t> buft_idx;
+    groups.clear();
+    for (size_t i = 0; i < n_tensors; i++) {
+        ggml_backend_buffer_t buf = get_buffer(i);
+        auto it = buft_idx.find(buf->buft);
+        if (it == buft_idx.end()) {
+            buft_idx[buf->buft] = groups.size();
+            groups.emplace_back();
+        }
+        groups[buft_idx[buf->buft]].push_back(i);
+    }
+}
+
+static bool ggml_backend_tensor_batch_is_2d(
+        const struct ggml_tensor * const * tensors,
+        const size_t * offsets,
+        const size_t * sizes,
+        const size_t * data_offsets,
+        size_t n_tensors,
+        size_t * base_offset,
+        size_t * size,
+        size_t * stride_tensor,
+        size_t * stride_data) {
+    if (n_tensors <= 1) {
+        return false;
+    }
+
+    const struct ggml_tensor * tensor = tensors[0];
+    const size_t size0 = sizes[0];
+
+    if (size0 == 0 || offsets[1] < offsets[0] || data_offsets[1] < data_offsets[0]) {
+        return false;
+    }
+
+    const size_t stride_tensor0 = offsets[1] - offsets[0];
+    const size_t stride_data0   = data_offsets[1] - data_offsets[0];
+
+    for (size_t i = 1; i < n_tensors; i++) {
+        if (tensors[i] != tensor || sizes[i] != size0) {
+            return false;
+        }
+        if (offsets[i] != offsets[0] + i*stride_tensor0 || data_offsets[i] != data_offsets[0] + i*stride_data0) {
+            return false;
+        }
+    }
+
+    *base_offset   = offsets[0];
+    *size          = size0;
+    *stride_tensor = stride_tensor0;
+    *stride_data   = stride_data0;
+
+    return true;
+}
+
 void ggml_backend_tensor_set_async(ggml_backend_t backend, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
     GGML_ASSERT(backend);
     GGML_ASSERT(tensor);
@@ -546,6 +602,154 @@ void ggml_backend_tensor_get_2d(const struct ggml_tensor * tensor, void * data,
     buf->iface.get_tensor_2d(buf, tensor, data, offset, size, n_copies, stride_tensor, stride_data);
 }
 
+void ggml_backend_tensor_set_n(struct ggml_tensor * const * tensors, const void * data,
+        const size_t * offsets, const size_t * sizes, const size_t * data_offsets, size_t n_tensors) {
+    if (n_tensors == 0) {
+        return;
+    }
+
+    std::vector<const void *> data_ptrs(n_tensors);
+    for (size_t i = 0; i < n_tensors; i++) {
+        data_ptrs[i] = (const char *) data + data_offsets[i];
+    }
+
+    ggml_backend_buffer_t buf = NULL;
+    bool same_buft = true;
+
+    for (size_t i = 0; i < n_tensors; i++) {
+        GGML_ASSERT(tensors[i] != NULL);
+        ggml_backend_buffer_t tensor_buf = tensors[i]->view_src ? tensors[i]->view_src->buffer : tensors[i]->buffer;
+        GGML_ASSERT(tensor_buf != NULL && "tensor buffer not set");
+        GGML_ASSERT(tensors[i]->data != NULL && "tensor not allocated");
+        GGML_ASSERT(offsets[i] + sizes[i] <= ggml_nbytes(tensors[i]) && "tensor write out of bounds");
+
+        if (buf == NULL) {
+            buf = tensor_buf;
+        } else if (buf->buft != tensor_buf->buft) {
+            same_buft = false;
+        }
+    }
+
+    if (same_buft && buf->iface.set_tensor_n != NULL) {
+        buf->iface.set_tensor_n(buf, tensors, data_ptrs.data(), offsets, sizes, n_tensors);
+        return;
+    }
+
+    if (!same_buft) {
+        std::vector<std::vector<size_t>> groups;
+        ggml_backend_tensor_buft_groups(n_tensors, [&](size_t i) {
+            return tensors[i]->view_src ? tensors[i]->view_src->buffer : tensors[i]->buffer;
+        }, groups);
+
+        for (size_t g = 0; g < groups.size(); g++) {
+            std::vector<ggml_tensor *> g_tensors;
+            std::vector<size_t> g_offsets;
+            std::vector<size_t> g_sizes;
+            std::vector<size_t> g_data_offsets;
+            g_tensors.reserve(groups[g].size());
+            g_offsets.reserve(groups[g].size());
+            g_sizes.reserve(groups[g].size());
+            g_data_offsets.reserve(groups[g].size());
+
+            for (size_t i : groups[g]) {
+                g_tensors.push_back(tensors[i]);
+                g_offsets.push_back(offsets[i]);
+                g_sizes.push_back(sizes[i]);
+                g_data_offsets.push_back(data_offsets[i]);
+            }
+
+            ggml_backend_tensor_set_n(g_tensors.data(), data, g_offsets.data(), g_sizes.data(), g_data_offsets.data(), g_tensors.size());
+        }
+        return;
+    }
+
+    size_t base_offset = 0, size = 0, stride_tensor = 0, stride_data = 0;
+    if (ggml_backend_tensor_batch_is_2d(tensors, offsets, sizes, data_offsets, n_tensors,
+                &base_offset, &size, &stride_tensor, &stride_data)) {
+        ggml_backend_tensor_set_2d(tensors[0], data_ptrs[0], base_offset, size,
+                n_tensors, stride_tensor, stride_data);
+        return;
+    }
+
+    for (size_t i = 0; i < n_tensors; i++) {
+        ggml_backend_tensor_set(tensors[i], data_ptrs[i], offsets[i], sizes[i]);
+    }
+}
+
+void ggml_backend_tensor_get_n(const struct ggml_tensor * const * tensors, void * data,
+        const size_t * offsets, const size_t * sizes, const size_t * data_offsets, size_t n_tensors) {
+    if (n_tensors == 0) {
+        return;
+    }
+
+    std::vector<void *> data_ptrs(n_tensors);
+    for (size_t i = 0; i < n_tensors; i++) {
+        data_ptrs[i] = (char *) data + data_offsets[i];
+    }
+
+    ggml_backend_buffer_t buf = NULL;
+    bool same_buft = true;
+
+    for (size_t i = 0; i < n_tensors; i++) {
+        GGML_ASSERT(tensors[i] != NULL);
+        ggml_backend_buffer_t tensor_buf = tensors[i]->view_src ? tensors[i]->view_src->buffer : tensors[i]->buffer;
+        GGML_ASSERT(tensor_buf != NULL && "tensor buffer not set");
+        GGML_ASSERT(tensors[i]->data != NULL && "tensor not allocated");
+        GGML_ASSERT(offsets[i] + sizes[i] <= ggml_nbytes(tensors[i]) && "tensor read out of bounds");
+
+        if (buf == NULL) {
+            buf = tensor_buf;
+        } else if (buf->buft != tensor_buf->buft) {
+            same_buft = false;
+        }
+    }
+
+    if (same_buft && buf->iface.get_tensor_n != NULL) {
+        buf->iface.get_tensor_n(buf, tensors, data_ptrs.data(), offsets, sizes, n_tensors);
+        return;
+    }
+
+    if (!same_buft) {
+        std::vector<std::vector<size_t>> groups;
+        ggml_backend_tensor_buft_groups(n_tensors, [&](size_t i) {
+            return tensors[i]->view_src ? tensors[i]->view_src->buffer : tensors[i]->buffer;
+        }, groups);
+
+        for (size_t g = 0; g < groups.size(); g++) {
+            std::vector<const ggml_tensor *> g_tensors;
+            std::vector<size_t> g_offsets;
+            std::vector<size_t> g_sizes;
+            std::vector<size_t> g_data_offsets;
+            g_tensors.reserve(groups[g].size());
+            g_offsets.reserve(groups[g].size());
+            g_sizes.reserve(groups[g].size());
+            g_data_offsets.reserve(groups[g].size());
+
+            for (size_t i : groups[g]) {
+                g_tensors.push_back(tensors[i]);
+                g_offsets.push_back(offsets[i]);
+                g_sizes.push_back(sizes[i]);
+                g_data_offsets.push_back(data_offsets[i]);
+            }
+
+            ggml_backend_tensor_get_n(g_tensors.data(), data, g_offsets.data(), g_sizes.data(), g_data_offsets.data(), g_tensors.size());
+        }
+        return;
+    }
+
+    size_t base_offset = 0, size = 0, stride_tensor = 0, stride_data = 0;
+    if (ggml_backend_tensor_batch_is_2d(tensors, offsets, sizes, data_offsets, n_tensors,
+                &base_offset, &size, &stride_tensor, &stride_data)) {
+        ggml_backend_tensor_get_2d(tensors[0], data_ptrs[0], base_offset, size,
+                n_tensors, stride_tensor, stride_data);
+        return;
+    }
+
+    for (size_t i = 0; i < n_tensors; i++) {
+        ggml_backend_tensor_get(tensors[i], data_ptrs[i], offsets[i], sizes[i]);
+    }
+}
+
 void ggml_backend_tensor_memset(struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) {
     GGML_ASSERT(tensor);
     ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
@@ -648,6 +852,59 @@ void ggml_backend_tensor_copy(const struct ggml_tensor * src, struct ggml_tensor
     }
 }
 
+void ggml_backend_tensor_copy_n(const struct ggml_tensor * const * srcs, struct ggml_tensor * const * dsts, size_t n_tensors) {
+    if (n_tensors == 0) {
+        return;
+    }
+
+    for (size_t i = 0; i < n_tensors; ++i) {
+        GGML_ASSERT(ggml_are_same_layout(srcs[i], dsts[i]) && "cannot copy tensors with different layouts");
+    }
+
+    ggml_backend_buffer_t buf = dsts[0]->view_src ? dsts[0]->view_src->buffer : dsts[0]->buffer;
+    bool same_buft = true;
+    for (size_t i = 0; i < n_tensors; ++i) {
+        ggml_backend_buffer_t dst_buf = dsts[i]->view_src ? dsts[i]->view_src->buffer : dsts[i]->buffer;
+        if (dst_buf->buft != buf->buft) {
+            same_buft = false;
+            break;
+        }
+    }
+
+    // if cpy_tensor_n returns false, it must not have performed any copies
+    if (same_buft && buf->iface.cpy_tensor_n != NULL) {
+        if (buf->iface.cpy_tensor_n(buf, srcs, dsts, n_tensors)) {
+            return;
+        }
+    }
+
+    if (!same_buft) {
+        std::vector<std::vector<size_t>> groups;
+        ggml_backend_tensor_buft_groups(n_tensors, [&](size_t i) {
+            return dsts[i]->view_src ? dsts[i]->view_src->buffer : dsts[i]->buffer;
+        }, groups);
+
+        for (size_t g = 0; g < groups.size(); g++) {
+            std::vector<const ggml_tensor *> g_srcs;
+            std::vector<ggml_tensor *> g_dsts;
+            g_srcs.reserve(groups[g].size());
+            g_dsts.reserve(groups[g].size());
+
+            for (size_t i : groups[g]) {
+                g_srcs.push_back(srcs[i]);
+                g_dsts.push_back(dsts[i]);
+            }
+
+            ggml_backend_tensor_copy_n(g_srcs.data(), g_dsts.data(), g_srcs.size());
+        }
+        return;
+    }
+
+    for (size_t i = 0; i < n_tensors; ++i) {
+        ggml_backend_tensor_copy(srcs[i], dsts[i]);
+    }
+}
+
 void ggml_backend_tensor_copy_async(ggml_backend_t backend_src, ggml_backend_t backend_dst, const struct ggml_tensor * src, struct ggml_tensor * dst) {
     GGML_ASSERT(ggml_are_same_layout(src, dst) && "cannot copy tensors with different layouts");
 
@@ -853,6 +1110,9 @@ static const struct ggml_backend_buffer_i ggml_backend_multi_buffer_i = {
     /* .cpy_tensor      = */ NULL,
     /* .clear           = */ ggml_backend_multi_buffer_clear,
     /* .reset           = */ NULL,
+    /* .set_tensor_n    = */ NULL,
+    /* .get_tensor_n    = */ NULL,
+    /* .cpy_tensor_n    = */ NULL,
 };
 
 ggml_backend_buffer_t ggml_backend_multi_buffer_alloc_buffer(ggml_backend_buffer_t * buffers, size_t n_buffers) {
@@ -2532,6 +2792,9 @@ static const struct ggml_backend_buffer_i ggml_backend_cpu_buffer_i = {
     /* .cpy_tensor      = */ ggml_backend_cpu_buffer_cpy_tensor,
     /* .clear           = */ ggml_backend_cpu_buffer_clear,
     /* .reset           = */ NULL,
+    /* .set_tensor_n    = */ NULL,
+    /* .get_tensor_n    = */ NULL,
+    /* .cpy_tensor_n    = */ NULL,
 };
 
 static const struct ggml_backend_buffer_i ggml_backend_cpu_buffer_from_ptr_i = {
@@ -2546,6 +2809,9 @@ static const struct ggml_backend_buffer_i ggml_backend_cpu_buffer_from_ptr_i = {
     /* .cpy_tensor      = */ ggml_backend_cpu_buffer_cpy_tensor,
     /* .clear           = */ ggml_backend_cpu_buffer_clear,
     /* .reset           = */ NULL,
+    /* .set_tensor_n    = */ NULL,
+    /* .get_tensor_n    = */ NULL,
+    /* .cpy_tensor_n    = */ NULL,
 };
 
 // CPU backend buffer type

ggml/src/ggml-cann/ggml-cann.cpp

@@ -1478,6 +1478,9 @@ static const ggml_backend_buffer_i ggml_backend_cann_buffer_interface = {
     /* .cpy_tensor      = */ ggml_backend_cann_buffer_cpy_tensor,
     /* .clear           = */ ggml_backend_cann_buffer_clear,
     /* .reset           = */ NULL,
+    /* .set_tensor_n    = */ NULL,
+    /* .get_tensor_n    = */ NULL,
+    /* .cpy_tensor_n    = */ NULL,
 };
 
 /**
@@ -2667,6 +2670,10 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
             return true;
         case GGML_OP_FLASH_ATTN_EXT:
             {
+                // K/V row lists (ggml_flash_attn_ext_rows) are not supported
+                if (op->src[5] != nullptr) {
+                    return false;
+                }
 #ifdef ASCEND_310P
                 // FA not support on 310p device
                 return false;

ggml/src/ggml-cpu/amx/amx.cpp

@@ -116,6 +116,9 @@ static ggml_backend_buffer_i ggml_backend_amx_buffer_interface = {
     /* .cpy_tensor      = */ nullptr,
     /* .clear           = */ ggml_backend_amx_buffer_clear,
     /* .reset           = */ nullptr,
+    /* .set_tensor_n    = */ NULL,
+    /* .get_tensor_n    = */ NULL,
+    /* .cpy_tensor_n    = */ NULL,
 };
 
 static const char * ggml_backend_amx_buffer_type_get_name(ggml_backend_buffer_type_t buft) {

ggml/src/ggml-cpu/common.h

@@ -43,6 +43,14 @@ static inline float f32_to_f32(float x) {
     return x;
 }
 
+static inline float f8_e4m3_to_f32(ggml_fp8_e4m3_t x) {
+    return ggml_f8_e4m3_to_fp32(x.bits);
+}
+
+static inline ggml_fp8_e4m3_t f32_to_f8_e4m3(float x) {
+    return { ggml_fp32_to_f8_e4m3(x) };
+}
+
 // TODO - merge this into the traits table, after using row-based conversions
 template <class T>
 struct type_conversion_table;
@@ -65,6 +73,12 @@ struct type_conversion_table<ggml_bf16_t> {
     static constexpr ggml_bf16_t (*from_f32)(float) = f32_to_bf16;
 };
 
+template <>
+struct type_conversion_table<ggml_fp8_e4m3_t> {
+    static constexpr float (*to_f32)(ggml_fp8_e4m3_t) = f8_e4m3_to_f32;
+    static constexpr ggml_fp8_e4m3_t (*from_f32)(float) = f32_to_f8_e4m3;
+};
+
 template <>
 struct type_conversion_table<int32_t> {
     static constexpr float (*to_f32)(int32_t) = i32_to_f32;

ggml/src/ggml-cpu/ggml-cpu-impl.h

@@ -30,7 +30,7 @@ struct ggml_compute_params {
 };
 
 
-#if defined(_MSC_VER)
+#if defined(_MSC_VER) && !defined(__clang__)
 
 #define m512bh(p) p
 #define m512i(p) p

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

@@ -86,6 +86,9 @@ float ggml_table_f32_e8m0_half[1 << 8];
 // precomputed f32 table for ue4m3 (1 KB) (simd-mappings.h)
 float ggml_table_f32_ue4m3[1 << 8];
 
+// Signed FP8 needs its own table: UE4M3 ignores the sign, halves values for NVFP4, and maps NaN to zero.
+float ggml_table_f32_f8_e4m3[1 << 8];
+
 #if defined(__ARM_ARCH)
 struct ggml_arm_arch_features_type {
     int sve_cnt;
@@ -296,6 +299,12 @@ static const struct ggml_type_traits_cpu type_traits_cpu[GGML_TYPE_COUNT] = {
         .vec_dot_type             = GGML_TYPE_Q8_0,
         .nrows                    = 1,
     },
+    [GGML_TYPE_F8_E4M3] = {
+        .from_float               = quantize_row_f8_e4m3,
+        .vec_dot                  = ggml_vec_dot_f8_e4m3_f32,
+        .vec_dot_type             = GGML_TYPE_F32,
+        .nrows                    = 1,
+    },
     [GGML_TYPE_Q2_K] = {
         .from_float               = quantize_row_q2_K,
         .vec_dot                  = ggml_vec_dot_q2_K_q8_K,
@@ -2876,7 +2885,7 @@ struct ggml_cplan ggml_graph_plan(
                 case GGML_OP_ADD_ID:
                 case GGML_OP_ADD1:
                     {
-                        if (ggml_is_quantized(node->src[0]->type)) {
+                        if (ggml_is_quantized(node->src[0]->type) || node->src[0]->type == GGML_TYPE_F8_E4M3) {
                             cur = ggml_type_size(GGML_TYPE_F32) * node->src[0]->ne[0] * n_tasks;
                         }
                     } break;
@@ -2924,7 +2933,7 @@ struct ggml_cplan ggml_graph_plan(
                     } break;
                 case GGML_OP_OUT_PROD:
                     {
-                        if (ggml_is_quantized(node->src[0]->type) ||
+                        if (ggml_is_quantized(node->src[0]->type) || node->src[0]->type == GGML_TYPE_F8_E4M3 ||
                             node->src[0]->type == GGML_TYPE_F16) {
                             cur = ggml_type_size(GGML_TYPE_F32) * node->src[0]->ne[0] * n_tasks;
                         }
@@ -3908,6 +3917,7 @@ void ggml_cpu_init(void) {
             // initialize UE4M3 table (256 entries)
             for (int i = 0; i < (1 << 8); ++i) {
                 ggml_table_f32_ue4m3[i] = ggml_ue4m3_to_fp32(i);
+                ggml_table_f32_f8_e4m3[i] = ggml_f8_e4m3_to_fp32(i);
             }
 
             const uint64_t t_end = ggml_time_us(); UNUSED(t_end);

ggml/src/ggml-cpu/ops.cpp

@@ -537,26 +537,36 @@ void ggml_compute_forward_dup(
     switch (src0->type) {
         case GGML_TYPE_F16:
             {
-                /**/ if (dst->type == GGML_TYPE_F16)  ggml_compute_forward_dup_flt<ggml_fp16_t, ggml_fp16_t>(params, dst);
-                else if (dst->type == GGML_TYPE_BF16) ggml_compute_forward_dup_flt<ggml_fp16_t, ggml_bf16_t>(params, dst);
-                else if (dst->type == GGML_TYPE_F32)  ggml_compute_forward_dup_flt<ggml_fp16_t, float      >(params, dst);
+                /**/ if (dst->type == GGML_TYPE_F16)     ggml_compute_forward_dup_flt<ggml_fp16_t, ggml_fp16_t>(params, dst);
+                else if (dst->type == GGML_TYPE_BF16)    ggml_compute_forward_dup_flt<ggml_fp16_t, ggml_bf16_t>(params, dst);
+                else if (dst->type == GGML_TYPE_F32)     ggml_compute_forward_dup_flt<ggml_fp16_t, float      >(params, dst);
+                else if (dst->type == GGML_TYPE_F8_E4M3) ggml_compute_forward_dup_flt<ggml_fp16_t, ggml_fp8_e4m3_t>(params, dst);
                 else ggml_compute_forward_dup_to_q<ggml_fp16_t>(params, dst);
             } break;
         case GGML_TYPE_BF16:
             {
-                /**/ if (dst->type == GGML_TYPE_F16)  ggml_compute_forward_dup_flt<ggml_bf16_t, ggml_fp16_t>(params, dst);
-                else if (dst->type == GGML_TYPE_BF16) ggml_compute_forward_dup_flt<ggml_bf16_t, ggml_bf16_t>(params, dst);
-                else if (dst->type == GGML_TYPE_F32)  ggml_compute_forward_dup_flt<ggml_bf16_t, float      >(params, dst);
+                /**/ if (dst->type == GGML_TYPE_F16)     ggml_compute_forward_dup_flt<ggml_bf16_t, ggml_fp16_t>(params, dst);
+                else if (dst->type == GGML_TYPE_BF16)    ggml_compute_forward_dup_flt<ggml_bf16_t, ggml_bf16_t>(params, dst);
+                else if (dst->type == GGML_TYPE_F32)     ggml_compute_forward_dup_flt<ggml_bf16_t, float      >(params, dst);
+                else if (dst->type == GGML_TYPE_F8_E4M3) ggml_compute_forward_dup_flt<ggml_bf16_t, ggml_fp8_e4m3_t>(params, dst);
                 else ggml_compute_forward_dup_to_q<ggml_bf16_t>(params, dst);
             } break;
         case GGML_TYPE_F32:
             {
-                /**/ if (dst->type == GGML_TYPE_F16)  ggml_compute_forward_dup_flt<float, ggml_fp16_t>(params, dst);
-                else if (dst->type == GGML_TYPE_BF16) ggml_compute_forward_dup_flt<float, ggml_bf16_t>(params, dst);
-                else if (dst->type == GGML_TYPE_F32)  ggml_compute_forward_dup_flt<float, float      >(params, dst);
-                else if (dst->type == GGML_TYPE_I32)  ggml_compute_forward_dup_flt<float, int32_t    >(params, dst);
+                /**/ if (dst->type == GGML_TYPE_F16)     ggml_compute_forward_dup_flt<float, ggml_fp16_t>(params, dst);
+                else if (dst->type == GGML_TYPE_BF16)    ggml_compute_forward_dup_flt<float, ggml_bf16_t>(params, dst);
+                else if (dst->type == GGML_TYPE_F32)     ggml_compute_forward_dup_flt<float, float      >(params, dst);
+                else if (dst->type == GGML_TYPE_F8_E4M3) ggml_compute_forward_dup_flt<float, ggml_fp8_e4m3_t>(params, dst);
+                else if (dst->type == GGML_TYPE_I32)     ggml_compute_forward_dup_flt<float, int32_t    >(params, dst);
                 else ggml_compute_forward_dup_to_q<float>(params, dst);
             } break;
+        case GGML_TYPE_F8_E4M3:
+            {
+                /**/ if (dst->type == GGML_TYPE_F16)  ggml_compute_forward_dup_flt<ggml_fp8_e4m3_t, ggml_fp16_t>(params, dst);
+                else if (dst->type == GGML_TYPE_BF16) ggml_compute_forward_dup_flt<ggml_fp8_e4m3_t, ggml_bf16_t>(params, dst);
+                else if (dst->type == GGML_TYPE_F32)  ggml_compute_forward_dup_flt<ggml_fp8_e4m3_t, float      >(params, dst);
+                else GGML_ABORT("not implemented");
+            } break;
         case GGML_TYPE_I32:
             {
                 if (dst->type == GGML_TYPE_F32) ggml_compute_forward_dup_flt<int32_t, float>(params, dst);
@@ -605,8 +615,9 @@ static void ggml_compute_forward_add_q_f32(
     GGML_ASSERT(nb1 <= nb2);
     GGML_ASSERT(nb2 <= nb3);
 
-    GGML_ASSERT(ggml_is_quantized(src0->type));
+    GGML_ASSERT(ggml_is_quantized(src0->type) || src0->type == GGML_TYPE_F8_E4M3);
     GGML_ASSERT(src1->type == GGML_TYPE_F32);
+    GGML_ASSERT(ne00 % ggml_blck_size(type) == 0);
 
     // rows per thread
     const int dr = (nr + nth - 1)/nth;
@@ -636,8 +647,6 @@ static void ggml_compute_forward_add_q_f32(
         float * src1_row = (float *)((char *) src1->data + (i11*nb11 + i12*nb12 + i13*nb13));
         void  * dst_row  = (void *) ((char *)  dst->data + ( i1*nb1  +  i2*nb2  +  i3*nb3));
 
-        assert(ne00 % 32 == 0);
-
         // unquantize row from src0 to temp buffer
         dequantize_row_q(src0_row, wdata, ne00);
         // add src1
@@ -673,6 +682,7 @@ void ggml_compute_forward_add(
         case GGML_TYPE_Q8_0:
         case GGML_TYPE_MXFP4:
         case GGML_TYPE_NVFP4:
+        case GGML_TYPE_F8_E4M3:
         case GGML_TYPE_Q2_K:
         case GGML_TYPE_Q3_K:
         case GGML_TYPE_Q4_K:
@@ -955,7 +965,7 @@ static void ggml_compute_forward_add1_q_f32(
     GGML_ASSERT(nb1 <= nb2);
     GGML_ASSERT(nb2 <= nb3);
 
-    GGML_ASSERT(ggml_is_quantized(src0->type));
+    GGML_ASSERT(ggml_is_quantized(src0->type) || src0->type == GGML_TYPE_F8_E4M3);
     GGML_ASSERT(dst->type == src0->type);
     GGML_ASSERT(src1->type == GGML_TYPE_F32);
 
@@ -1130,6 +1140,7 @@ void ggml_compute_forward_add1(
         case GGML_TYPE_Q8_1:
         case GGML_TYPE_MXFP4:
         case GGML_TYPE_NVFP4:
+        case GGML_TYPE_F8_E4M3:
         case GGML_TYPE_Q2_K:
         case GGML_TYPE_Q3_K:
         case GGML_TYPE_Q4_K:
@@ -1261,6 +1272,7 @@ void ggml_compute_forward_acc(
         case GGML_TYPE_Q8_1:
         case GGML_TYPE_MXFP4:
         case GGML_TYPE_NVFP4:
+        case GGML_TYPE_F8_E4M3:
         case GGML_TYPE_Q2_K:
         case GGML_TYPE_Q3_K:
         case GGML_TYPE_Q4_K:
@@ -4804,6 +4816,7 @@ void ggml_compute_forward_out_prod(
         case GGML_TYPE_Q8_0:
         case GGML_TYPE_MXFP4:
         case GGML_TYPE_NVFP4:
+        case GGML_TYPE_F8_E4M3:
         case GGML_TYPE_Q2_K:
         case GGML_TYPE_Q3_K:
         case GGML_TYPE_Q4_K:
@@ -5129,6 +5142,7 @@ void ggml_compute_forward_set(
         case GGML_TYPE_Q8_1:
         case GGML_TYPE_MXFP4:
         case GGML_TYPE_NVFP4:
+        case GGML_TYPE_F8_E4M3:
         case GGML_TYPE_Q2_K:
         case GGML_TYPE_Q3_K:
         case GGML_TYPE_Q4_K:
@@ -5354,6 +5368,7 @@ void ggml_compute_forward_get_rows(
         case GGML_TYPE_Q8_1:
         case GGML_TYPE_MXFP4:
         case GGML_TYPE_NVFP4:
+        case GGML_TYPE_F8_E4M3:
         case GGML_TYPE_Q2_K:
         case GGML_TYPE_Q3_K:
         case GGML_TYPE_Q4_K:
@@ -6114,6 +6129,7 @@ void ggml_compute_forward_clamp(
         case GGML_TYPE_Q8_1:
         case GGML_TYPE_MXFP4:
         case GGML_TYPE_NVFP4:
+        case GGML_TYPE_F8_E4M3:
         case GGML_TYPE_Q2_K:
         case GGML_TYPE_Q3_K:
         case GGML_TYPE_Q4_K:
@@ -8828,6 +8844,7 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk(
     const ggml_tensor * v     = dst->src[2];
     const ggml_tensor * mask  = dst->src[3];
     const ggml_tensor * sinks = dst->src[4];
+    const ggml_tensor * kv_rows = dst->src[5];
 
     GGML_TENSOR_LOCALS(int64_t, neq, q,   ne)
     GGML_TENSOR_LOCALS(size_t,  nbq, q,   nb)
@@ -8935,6 +8952,8 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk(
         const float * pq = (const float *) ((char *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3));
         q_to_vec_dot(pq, Q_q, DK);
 
+        const int32_t * rows = kv_rows ? (const int32_t *) ((const char *) kv_rows->data + iq3*kv_rows->nb[1]) : nullptr;
+
         // online softmax / attention
         // loop over n_kv and n_head_kv
         // ref: https://arxiv.org/pdf/2112.05682.pdf
@@ -8945,9 +8964,14 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk(
                 continue;
             }
 
+            const int64_t ik1 = rows ? rows[ic] : ic;
+            if (ik1 < 0) {
+                continue;
+            }
+
             float s; // KQ value
 
-            const char * k_data = (const char *) k->data + ( ic*nbk1 + ik2*nbk2 + ik3*nbk3);
+            const char * k_data = (const char *) k->data + (ik1*nbk1 + ik2*nbk2 + ik3*nbk3);
             kq_vec_dot(DK, &s, 0, k_data, 0, Q_q, 0, 1);
 
             s = s*scale; // scale KQ value
@@ -8963,7 +8987,7 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk(
             float ms = 1.0f; // upon new higher max val, scale VKQ and KQ sum with this value
             float vs = 1.0f; // post-softmax KQ value, expf(s - M)
 
-            const char * v_data = ((const char *) v->data + (ic*nbv1 + iv2*nbv2 + iv3*nbv3));
+            const char * v_data = ((const char *) v->data + (ik1*nbv1 + iv2*nbv2 + iv3*nbv3));
 
             if (v->type == GGML_TYPE_F16) {
                 if (s > M) {
@@ -9130,6 +9154,10 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
     static constexpr int Q_TILE_SZ  = ggml_fa_tile_config::Q;
     static constexpr int KV_TILE_SZ = ggml_fa_tile_config::KV;
 
+    // with kv_rows the loop runs over the mask columns and reads K/V rows through kv_rows
+    const ggml_tensor * kv_rows = dst->src[5];
+    const int64_t n_kv_cols = kv_rows ? kv_rows->ne[0] : nek1;
+
     int ir = ir0;
     while (ir < ir1) {
         // q indices for the start of this tile
@@ -9137,6 +9165,10 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
         const int iq2 = (ir - iq3*neq2*neq1)/neq1;
         const int iq1 = (ir - iq3*neq2*neq1 - iq2*neq1);
 
+        // padding rows are masked, read row 0 for them
+        const int32_t * rows = kv_rows ? (const int32_t *) ((const char *) kv_rows->data + iq3*kv_rows->nb[1]) : nullptr;
+        auto kv_row = [rows](int64_t i) -> int64_t { return rows ? std::max<int32_t>(rows[i], 0) : i; };
+
         // Number of valid rows in this tile:
         // - limited by tile size (Q_TILE_SZ)
         // - limited by chunk boundary (ir1 - ir)
@@ -9196,8 +9228,8 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
         memset(K_f32, 0, DK * KV_TILE_SZ * sizeof(float));
         memset(V32,   0, KV_TILE_SZ * DV * sizeof(float));
 
-        for (int64_t ic = 0; ic < nek1; ic += KV_TILE_SZ) {
-            const int kv_tile = (int)std::min((int64_t)KV_TILE_SZ, nek1 - ic);
+        for (int64_t ic = 0; ic < n_kv_cols; ic += KV_TILE_SZ) {
+            const int kv_tile = (int)std::min((int64_t)KV_TILE_SZ, n_kv_cols - ic);
 
             // skip the tile entirely if all the masks are -inf
             if (mask) {
@@ -9224,7 +9256,7 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
             // Pack K tile transposed: K_f32[dk][kv] so KV_TILE is contiguous (SIMD dim)
             // Zero-pad the last tile so the GEMM always operates on KV_TILE_SZ columns
             for (int tk = 0; tk < kv_tile; tk++) {
-                const char * k_data = (const char *)k->data + (ic + tk)*nbk1 + ik2*nbk2 + ik3*nbk3;
+                const char * k_data = (const char *)k->data + kv_row(ic + tk)*nbk1 + ik2*nbk2 + ik3*nbk3;
                 if (kv_type == GGML_TYPE_F16) {
                     const ggml_fp16_t * k_f16 = (const ggml_fp16_t *)k_data;
                     for (int64_t dk = 0; dk < DK; dk++) {
@@ -9289,7 +9321,7 @@ static void ggml_compute_forward_flash_attn_ext_tiled(
             // V accumulation: VKQ32 += softmax(KQ) * V
             // Pack V tile to contiguous F32, zero-padded
             for (int tk = 0; tk < kv_tile; tk++) {
-                const char * v_data = (const char *)v->data + (ic + tk)*nbv1 + iv2*nbv2 + iv3*nbv3;
+                const char * v_data = (const char *)v->data + kv_row(ic + tk)*nbv1 + iv2*nbv2 + iv3*nbv3;
                 if (kv_type == GGML_TYPE_F16) {
                     ggml_cpu_fp16_to_fp32((const ggml_fp16_t *)v_data, V32 + tk * DV, DV);
                 } else {
@@ -9461,8 +9493,12 @@ static void ggml_compute_forward_flash_attn_ext_f16(
     // When use_ref is set, force the vec-only reference implementation (no tiling, no KV-chunking)
     const bool use_ref = params->use_ref;
 
+    // with kv_rows the loops run over the mask columns, the split KV path does not support it
+    const ggml_tensor * kv_rows = dst->src[5];
+    const int64_t n_kv_cols = kv_rows ? kv_rows->ne[0] : nek1;
+
     const bool kv_is_f32_or_f16 = (k->type == GGML_TYPE_F32 || k->type == GGML_TYPE_F16);
-    const bool use_split_kv_path = !use_ref && (neq1 == 1 && neq3 == 1) && kv_is_f32_or_f16 && (k->type == v->type) && q->type == GGML_TYPE_F32 && nek1 >= 512;
+    const bool use_split_kv_path = !use_ref && !kv_rows && (neq1 == 1 && neq3 == 1) && kv_is_f32_or_f16 && (k->type == v->type) && q->type == GGML_TYPE_F32 && nek1 >= 512;
 
     if (use_split_kv_path) {
         const int64_t chunk_size = (nek1 + nth - 1) / nth;
@@ -9541,7 +9577,7 @@ static void ggml_compute_forward_flash_attn_ext_f16(
             if (use_tiled) {
                 ggml_compute_forward_flash_attn_ext_tiled(params, dst, ir0, ir1);
             } else {
-                ggml_compute_forward_flash_attn_ext_f16_one_chunk(params, dst, ir0, ir1, 0, nek1, nullptr, 0);
+                ggml_compute_forward_flash_attn_ext_f16_one_chunk(params, dst, ir0, ir1, 0, n_kv_cols, nullptr, 0);
             }
 
             current_chunk = ggml_threadpool_chunk_add(params->threadpool, 1);

ggml/src/ggml-cpu/quants.c

@@ -3,6 +3,7 @@
 
 #include "ggml-cpu-impl.h"
 #include "simd-mappings.h"
+#include "vec.h"
 #include "ggml-quants.h"
 #include "quants.h"
 
@@ -62,6 +63,10 @@ void quantize_row_nvfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, i
     quantize_row_nvfp4_ref(x, y, k);
 }
 
+void quantize_row_f8_e4m3(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) {
+    quantize_row_f8_e4m3_ref(x, y, k);
+}
+
 //
 // 2-6 bit quantization in super-blocks
 //
@@ -362,6 +367,49 @@ void ggml_vec_dot_nvfp4_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs,
     *s = sumf;
 }
 
+void ggml_vec_dot_f8_e4m3_f32(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const ggml_fp8_e4m3_t * GGML_RESTRICT x = vx;
+    const float * GGML_RESTRICT y = vy;
+
+    float sumf = 0.0f;
+    int pos = 0;
+#if defined(GGML_SIMD) && !defined(__ARM_FEATURE_SVE) && !defined(__riscv_v_intrinsic)
+    GGML_F32_VEC sum[GGML_F32_ARR] = { GGML_F32_VEC_ZERO };
+    const int np = n - n % GGML_F32_STEP;
+    for (; pos < np; pos += GGML_F32_STEP) {
+        for (int j = 0; j < GGML_F32_ARR; ++j) {
+            const int offset = pos + j*GGML_F32_EPR;
+            GGML_F32_VEC ax;
+#if defined(__AVX512F__) && defined(__AVX512BW__)
+            const __m128i codes = _mm_loadu_si128((const __m128i *) (x + offset));
+            ax = _mm512_i32gather_ps(_mm512_cvtepu8_epi32(codes), ggml_table_f32_f8_e4m3, sizeof(float));
+#elif defined(__AVX2__) && !defined(__AVX512F__)
+            const __m128i codes = _mm_loadl_epi64((const __m128i *) (x + offset));
+            ax = _mm256_i32gather_ps(ggml_table_f32_f8_e4m3, _mm256_cvtepu8_epi32(codes), sizeof(float));
+#else
+            float values[GGML_F32_EPR];
+            for (int k = 0; k < GGML_F32_EPR; ++k) {
+                values[k] = ggml_table_f32_f8_e4m3[x[offset + k].bits];
+            }
+            ax = GGML_F32_VEC_LOAD(values);
+#endif
+            sum[j] = GGML_F32_VEC_FMA(sum[j], ax, GGML_F32_VEC_LOAD(y + offset));
+        }
+    }
+    GGML_F32_VEC_REDUCE(sumf, sum);
+#endif
+    for (; pos < n; ++pos) {
+        sumf += ggml_table_f32_f8_e4m3[x[pos].bits] * y[pos];
+    }
+    *s = sumf;
+}
+
 void ggml_vec_dot_q5_0_q8_0_generic(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
     const int qk = QK8_0;
     const int nb = n / qk;

ggml/src/ggml-cpu/quants.h

@@ -23,6 +23,7 @@ void quantize_row_q8_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, in
 
 void quantize_row_mxfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
 void quantize_row_nvfp4(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
+void quantize_row_f8_e4m3(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
 
 void quantize_row_q2_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
 void quantize_row_q3_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k);
@@ -40,6 +41,7 @@ void quantize_row_iq4_xs (const float * GGML_RESTRICT x, void * GGML_RESTRICT y,
 // Dot product
 void ggml_vec_dot_q1_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
 void ggml_vec_dot_q2_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
+void ggml_vec_dot_f8_e4m3_f32(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
 void ggml_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
 void ggml_vec_dot_q4_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);
 void ggml_vec_dot_q5_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc);

ggml/src/ggml-cpu/simd-mappings.h

@@ -127,6 +127,8 @@ extern float ggml_table_f32_e8m0_half[1 << 8];
 // defined in ggml-cpu.c, initialized in ggml_cpu_init()
 extern float ggml_table_f32_ue4m3[1 << 8];
 
+extern float ggml_table_f32_f8_e4m3[1 << 8];
+
 // Use lookup table for E8M0 on x86 (faster than bit manipulation)
 #if defined(__AVX__) || defined(__AVX2__) || defined(__AVX512F__)
 #define GGML_CPU_E8M0_TO_FP32_HALF(x) ggml_table_f32_e8m0_half[(uint8_t)(x)]

ggml/src/ggml-cpu/spacemit/ime.cpp

@@ -1046,6 +1046,10 @@ class tensor_traits_common : public tensor_traits_base {
                         return true;
                 }
             case GGML_OP_FLASH_ATTN_EXT:
+                // K/V row lists are handled by the generic CPU path
+                if (op->src[5] != nullptr) {
+                    return false;
+                }
                 forward_flash_attn_ext_f16(params, op);
                 return true;
             case GGML_OP_CONT:
@@ -1489,6 +1493,9 @@ static const ggml_backend_buffer_i ggml_backend_riscv64_spacemit_buffer_i = {
     /* .cpy_tensor      = */ nullptr,
     /* .clear           = */ ggml_backend_riscv64_spacemit_buffer_clear,
     /* .reset           = */ nullptr,
+    /* .set_tensor_n    = */ NULL,
+    /* .get_tensor_n    = */ NULL,
+    /* .cpy_tensor_n    = */ NULL,
 };
 
 static const char * ggml_backend_cpu_riscv64_spacemit_buffer_type_get_name(ggml_backend_buffer_type_t buft) {

ggml/src/ggml-cuda/batch-get.cu

@@ -0,0 +1,198 @@
+#include "batch-get.cuh"
+#include "common.cuh"
+#include "ggml-cpp.h"
+
+#include <climits>
+
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+struct ggml_cuda_batch_get_descriptor {
+    const unsigned char * src;
+    unsigned char * dst;
+    size_t size;
+};
+
+static __global__ void ggml_cuda_batch_get_kernel(const ggml_cuda_batch_get_descriptor * copies, size_t count) {
+    // Blocks along y cooperate on each descriptor.
+    const size_t index = blockIdx.x;
+    if (index >= count) {
+        return;
+    }
+    const size_t first = (size_t) blockIdx.y * blockDim.x + threadIdx.x;
+    const size_t stride = (size_t) gridDim.y * blockDim.x;
+    const auto copy = copies[index];
+    // Use 16-, 8-, 4-, or 2-byte copies when both pointers meet that alignment; threads copy separate chunks.
+    const uintptr_t alignment = (uintptr_t) copy.src | (uintptr_t) copy.dst;
+    size_t tail = 0;
+    if (alignment % 16 == 0) {
+        for (size_t i = first; i < copy.size / 16; i += stride) {
+            ggml_cuda_memcpy_1<16>(copy.dst + 16*i, copy.src + 16*i);
+        }
+        tail = copy.size / 16 * 16;
+    } else if (alignment % 8 == 0) {
+        for (size_t i = first; i < copy.size / 8; i += stride) {
+            ggml_cuda_memcpy_1<8>(copy.dst + 8*i, copy.src + 8*i);
+        }
+        tail = copy.size / 8 * 8;
+    } else if (alignment % 4 == 0) {
+        for (size_t i = first; i < copy.size / 4; i += stride) {
+            ggml_cuda_memcpy_1<4>(copy.dst + 4*i, copy.src + 4*i);
+        }
+        tail = copy.size / 4 * 4;
+    } else if (alignment % 2 == 0) {
+        for (size_t i = first; i < copy.size / 2; i += stride) {
+            ggml_cuda_memcpy_1<2>(copy.dst + 2*i, copy.src + 2*i);
+        }
+        tail = copy.size / 2 * 2;
+    }
+    // Copy trailing bytes, or the entire range when no wider alignment is available.
+    for (size_t i = tail + first; i < copy.size; i += stride) {
+        ggml_cuda_memcpy_1<1>(copy.dst + i, copy.src + i);
+    }
+}
+
+struct ggml_cuda_batch_get_mapping {
+    void * host = nullptr;
+    unsigned char * device = nullptr;
+
+    ggml_cuda_batch_get_mapping(void * host, size_t size) {
+        const cudaError_t err = cudaHostRegister(host, size, cudaHostRegisterMapped | cudaHostRegisterPortable);
+        if (err != cudaSuccess) {
+            (void) cudaGetLastError();
+            GGML_LOG_DEBUG("%s: failed to register %.2f MiB of pinned memory: %s\n", __func__,
+                           size / 1024.0 / 1024.0, cudaGetErrorString(err));
+            return;
+        }
+        this->host = host;
+        if (cudaHostGetDevicePointer((void **) &device, host, 0) != cudaSuccess) {
+            (void) cudaGetLastError();
+            device = nullptr;
+            return;
+        }
+    }
+
+    ggml_cuda_batch_get_mapping(const ggml_cuda_batch_get_mapping &) = delete;
+    ggml_cuda_batch_get_mapping & operator=(const ggml_cuda_batch_get_mapping &) = delete;
+
+    ~ggml_cuda_batch_get_mapping() {
+        if (host) {
+            // Release only the registration acquired by this object.
+            CUDA_CHECK(cudaHostUnregister(host));
+        }
+    }
+};
+#endif
+
+bool ggml_cuda_get_tensor_n(int device,
+        const ggml_tensor * const * tensors, void * const * data, const size_t * offsets,
+        const size_t * sizes, size_t n_tensors) {
+#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
+    GGML_UNUSED(device);
+    GGML_UNUSED(tensors);
+    GGML_UNUSED(data);
+    GGML_UNUSED(offsets);
+    GGML_UNUSED(sizes);
+    GGML_UNUSED(n_tensors);
+    return false;
+#else
+    // Provisional dispatch values based on experiments.
+    constexpr size_t min_batch_reads = 128;
+    constexpr size_t max_gather_read_bytes = 128 * 1024;
+    // Collect non-empty CPU destination ranges.
+    std::vector<std::pair<uintptr_t, uintptr_t>> ranges;
+    ranges.reserve(n_tensors);
+    size_t payload_bytes = 0;
+    // Number of non-empty reads whose size is at most max_gather_read_bytes.
+    size_t n_reads_small = 0;
+    for (size_t i = 0; i < n_tensors; ++i) {
+        if (sizes[i] == 0) {
+            continue;
+        }
+        const uintptr_t address = (uintptr_t) data[i];
+        GGML_ASSERT(address <= UINTPTR_MAX - sizes[i]);
+        GGML_ASSERT(payload_bytes <= SIZE_MAX - sizes[i]);
+        payload_bytes += sizes[i];
+        n_reads_small += sizes[i] <= max_gather_read_bytes;
+        ranges.emplace_back(address, address + sizes[i]);
+    }
+    if (ranges.empty()) {
+        return true;
+    }
+    // Total number of non-empty reads.
+    const size_t n_reads = ranges.size();
+    if (n_reads > SIZE_MAX / sizeof(ggml_cuda_batch_get_descriptor) || n_reads > INT_MAX) {
+        return false;
+    }
+    if (n_reads < 2 || (n_reads < min_batch_reads && n_reads_small == n_reads)) {
+        return false;
+    }
+    std::sort(ranges.begin(), ranges.end());
+    for (size_t i = 1; i < ranges.size(); ++i) {
+        // Preserve the ordered copy path when CPU destinations overlap.
+        if (ranges[i].first < ranges[i - 1].second) {
+            return false;
+        }
+    }
+    // Find the CPU span to register, including gaps between destinations.
+    const uintptr_t begin = ranges.front().first;
+    const uintptr_t end = ranges.back().second;
+    if (end - begin - payload_bytes > payload_bytes) {
+        return false;
+    }
+    // Register and map the existing CPU destination so the kernel can write to it directly.
+    ggml_cuda_batch_get_mapping mapping((void *) begin, end - begin);
+    if (!mapping.host) {
+        return false;
+    }
+    // Reads selected for gather; zero means all reads use DMA.
+    size_t n_reads_gather = mapping.device && n_reads_small >= min_batch_reads ? n_reads_small : 0;
+    // Keep descriptor storage alive until the final stream synchronization.
+    constexpr size_t max_batch_descriptors = 128 * 1024;
+    const size_t bytes = std::min(n_reads_gather, max_batch_descriptors) * sizeof(ggml_cuda_batch_get_descriptor);
+    ggml_backend_buffer_ptr descriptors(n_reads_gather ? ggml_backend_buft_alloc_buffer(ggml_backend_cuda_buffer_type(device), bytes) : nullptr);
+    if (!descriptors) {
+        n_reads_gather = 0;
+    }
+    const auto host_type = ggml_backend_cuda_host_buffer_type();
+    ggml_backend_buffer_ptr host_descriptors(n_reads_gather ? ggml_backend_buft_alloc_buffer(host_type, bytes) : nullptr);
+    if (!host_descriptors || ggml_backend_buffer_get_type(host_descriptors.get()) != host_type) {
+        n_reads_gather = 0;
+    }
+    if (n_reads_gather) {
+        auto * copies = static_cast<ggml_cuda_batch_get_descriptor *>(ggml_backend_buffer_get_base(host_descriptors.get()));
+        auto * device_copies = static_cast<ggml_cuda_batch_get_descriptor *>(ggml_backend_buffer_get_base(descriptors.get()));
+        const int threads = 256;
+        const size_t nsm = ggml_cuda_info().devices[device].nsm;
+        size_t i = 0;
+        for (size_t first = 0; first < n_reads_gather; first += max_batch_descriptors) {
+            const size_t n_batch = std::min(n_reads_gather - first, max_batch_descriptors);
+            size_t count = 0;
+            for (; i < n_tensors && count < n_batch; ++i) {
+                if (sizes[i] == 0 || sizes[i] > max_gather_read_bytes) {
+                    continue;
+                }
+                copies[count++] = {(const unsigned char *) tensors[i]->data + offsets[i],
+                        mapping.device + ((uintptr_t) data[i] - begin), sizes[i]};
+            }
+            const size_t blocks_per_copy = std::min<size_t>(UINT16_MAX, std::max<size_t>(1, (nsm + n_batch - 1) / n_batch));
+            const dim3 blocks((unsigned int) n_batch, (unsigned int) blocks_per_copy);
+            // Stream ordering keeps each kernel ahead of the next descriptor overwrite.
+            CUDA_CHECK(cudaMemcpyAsync(device_copies, copies, n_batch * sizeof(copies[0]), cudaMemcpyHostToDevice, cudaStreamPerThread));
+            ggml_cuda_batch_get_kernel<<<blocks, threads, 0, cudaStreamPerThread>>>(device_copies, n_batch);
+            CUDA_CHECK(cudaGetLastError());
+            if (first + n_batch < n_reads_gather) {
+                // Wait before refilling the pinned host descriptor buffer.
+                CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
+            }
+        }
+    }
+    // Keep the destination pinned for DMA too; overlapping destinations already declined above.
+    for (size_t i = 0; i < n_tensors; ++i) {
+        if (sizes[i] && (!n_reads_gather || sizes[i] > max_gather_read_bytes)) {
+            CUDA_CHECK(cudaMemcpyAsync(data[i], (const char *) tensors[i]->data + offsets[i], sizes[i], cudaMemcpyDeviceToHost, cudaStreamPerThread));
+        }
+    }
+    // Wait for GPU writes before releasing the CPU mapping and descriptor buffers.
+    CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
+    return true;
+#endif
+}

ggml/src/ggml-cuda/batch-get.cuh

@@ -0,0 +1,7 @@
+#pragma once
+
+#include "ggml-backend.h"
+
+bool ggml_cuda_get_tensor_n(int device,
+        const ggml_tensor * const * tensors, void * const * data, const size_t * offsets,
+        const size_t * sizes, size_t n_tensors);

ggml/src/ggml-cuda/common.cuh

@@ -82,6 +82,7 @@
 #define GGML_CUDA_CC_RDNA2      (GGML_CUDA_CC_OFFSET_AMD + 0x1030) // RX 6000, minimum for dp4a
 #define GGML_CUDA_CC_RDNA3      (GGML_CUDA_CC_OFFSET_AMD + 0x1100) // RX 7000, minimum for WMMA
 #define GGML_CUDA_CC_RDNA3_5    (GGML_CUDA_CC_OFFSET_AMD + 0x1150) // AI 370, AI Max 395 laptops.
+#define GGML_CUDA_CC_RDNA4_M    (GGML_CUDA_CC_OFFSET_AMD + 0x1170) // gfx117X
 #define GGML_CUDA_CC_RDNA4      (GGML_CUDA_CC_OFFSET_AMD + 0x1200) // RX 9000
 
 #define GGML_CUDA_CC_IS_AMD(cc)     (cc >= GGML_CUDA_CC_OFFSET_AMD)
@@ -89,8 +90,9 @@
 #define GGML_CUDA_CC_IS_RDNA1(cc)   (cc >= GGML_CUDA_CC_RDNA1 && cc < GGML_CUDA_CC_RDNA2)
 #define GGML_CUDA_CC_IS_RDNA2(cc)   (cc >= GGML_CUDA_CC_RDNA2 && cc < GGML_CUDA_CC_RDNA3)
 #define GGML_CUDA_CC_IS_RDNA3_0(cc) (cc >= GGML_CUDA_CC_RDNA3 && cc < GGML_CUDA_CC_RDNA3_5)
-#define GGML_CUDA_CC_IS_RDNA3_5(cc) (cc >= GGML_CUDA_CC_RDNA3_5 && cc < GGML_CUDA_CC_RDNA4)
+#define GGML_CUDA_CC_IS_RDNA3_5(cc) (cc >= GGML_CUDA_CC_RDNA3_5 && cc < GGML_CUDA_CC_RDNA4_M)
 #define GGML_CUDA_CC_IS_RDNA3(cc)   (GGML_CUDA_CC_IS_RDNA3_0(cc) || GGML_CUDA_CC_IS_RDNA3_5(cc))
+#define GGML_CUDA_CC_IS_RDNA4_M(cc) (cc >= GGML_CUDA_CC_RDNA4_M && cc < GGML_CUDA_CC_RDNA4)
 #define GGML_CUDA_CC_IS_RDNA4(cc)   (cc >= GGML_CUDA_CC_RDNA4)
 #define GGML_CUDA_CC_IS_GCN_APU(cc) ((cc) == GGML_CUDA_CC_GFX909 || (cc) == GGML_CUDA_CC_GFX90C)
 #define GGML_CUDA_CC_IS_GCN(cc)     ((cc > GGML_CUDA_CC_OFFSET_AMD && cc < GGML_CUDA_CC_CDNA1) || GGML_CUDA_CC_IS_GCN_APU(cc))
@@ -277,9 +279,17 @@ static const char * cu_get_error_str(CUresult err) {
 #define AMD_MFMA_AVAILABLE
 #endif // defined(GGML_USE_HIP) && defined(CDNA) && !defined(GGML_HIP_NO_MMQ_MFMA)
 
-#if defined(GGML_USE_HIP) && (defined(RDNA4) || defined(RDNA3))
+#if defined(GGML_USE_HIP) && (defined(RDNA4) || defined(RDNA4_M))
+#define AMD_WMMA_128B_AVAILABLE
+#endif // defined(GGML_USE_HIP) && (defined(RDNA4) || defined(RDNA4_M))
+
+#if defined(GGML_USE_HIP) && defined(RDNA3)
+#define AMD_WMMA_256B_AVAILABLE
+#endif // defined(GGML_USE_HIP) && defined(RDNA3)
+
+#if defined(AMD_WMMA_128B_AVAILABLE) || defined(AMD_WMMA_256B_AVAILABLE)
 #define AMD_WMMA_AVAILABLE
-#endif // defined(GGML_USE_HIP) && defined(RDNA4)
+#endif // defined(AMD_WMMA_128B_AVAILABLE) || defined(AMD_WMMA_256B_AVAILABLE)
 
 // The Volta instructions are in principle available on Turing or newer but they are effectively unusable:
 #if !defined(GGML_USE_HIP) && __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA
@@ -326,7 +336,7 @@ static bool fast_fp16_hardware_available(const int cc) {
 // To be used for feature selection of external libraries, e.g. cuBLAS.
 static bool fp16_mma_hardware_available(const int cc) {
     return (GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_VOLTA) ||
-        GGML_CUDA_CC_IS_CDNA(cc) || GGML_CUDA_CC_IS_RDNA3(cc) || GGML_CUDA_CC_IS_RDNA4(cc) ||
+        GGML_CUDA_CC_IS_CDNA(cc) || GGML_CUDA_CC_IS_RDNA3(cc) || GGML_CUDA_CC_IS_RDNA4_M(cc) || GGML_CUDA_CC_IS_RDNA4(cc) ||
         (GGML_CUDA_CC_IS_MTHREADS(cc) && cc >= GGML_CUDA_CC_QY2);
 }
 
@@ -355,7 +365,7 @@ static bool amd_mfma_available(const int cc) {
 }
 
 static bool amd_wmma_available(const int cc) {
-    return (GGML_CUDA_CC_IS_RDNA4(cc) || GGML_CUDA_CC_IS_RDNA3(cc));
+    return GGML_CUDA_CC_IS_RDNA3(cc) || GGML_CUDA_CC_IS_RDNA4_M(cc) || GGML_CUDA_CC_IS_RDNA4(cc);
 }
 
 static bool volta_mma_available(const int cc) {
@@ -713,7 +723,7 @@ static __device__ __forceinline__ int ggml_cuda_dp4a(const int a, const int b, i
 #if defined(GGML_USE_HIP)
 #if defined(CDNA) || defined(RDNA2) || defined(__gfx906__)
     c = __builtin_amdgcn_sdot4(a, b, c, false);
-#elif defined(RDNA3) || defined(RDNA4)
+#elif defined(RDNA3) || defined(RDNA4_M) || defined(RDNA4)
     c = __builtin_amdgcn_sudot4( true, a, true, b, c, false);
 #elif defined(RDNA1) || defined(__gfx900__)
     int tmp1;
@@ -758,9 +768,9 @@ static __device__ __forceinline__ void ggml_cuda_mad(float & acc, const float2 v
     acc += v.y*u.y;
 }
 
-#if defined(GGML_USE_HIP) && (defined(RDNA2) || defined(RDNA3) || defined(RDNA4) || defined(__gfx906__) || defined(CDNA))
+#if defined(GGML_USE_HIP) && (defined(RDNA2) || defined(RDNA3) || defined(RDNA4_M) || defined(RDNA4) || defined(__gfx906__) || defined(CDNA))
 #define V_DOT2_F32_F16_AVAILABLE
-#endif // defined(GGML_USE_HIP) && (defined(RDNA2) || defined(RDNA3) || defined(RDNA4) || defined(__gfx906__) || defined(CDNA))
+#endif // defined(GGML_USE_HIP) && (defined(RDNA2) || defined(RDNA3) || defined(RDNA4_M) || defined(RDNA4) || defined(__gfx906__) || defined(CDNA))
 
 static __device__ __forceinline__ void ggml_cuda_mad(float & acc, const half2 v, const half2 u) {
 #ifdef V_DOT2_F32_F16_AVAILABLE
@@ -1459,16 +1469,21 @@ struct ggml_backend_cuda_context {
     // when the computation is split across CPU/GPU (e.g., with --n-cpu-moe)
     std::unordered_map<const void *, std::unique_ptr<ggml_cuda_graph>> cuda_graphs;
 
-    int64_t last_graph_eviction_sweep = 0;
+    int64_t next_graph_eviction_time = 0;
+    int64_t last_graph_use_time = 0;
 
-    ggml_cuda_graph * cuda_graph(const void * first_node_ptr) {
-        const int64_t time_now = ggml_time_us();
+    ggml_cuda_graph * cuda_graph(const void * first_node_ptr, int64_t time_now = ggml_time_us()) {
+        // Give all cached graphs 10s to be reused after an idle period, without updating each graph.
+        if (!cuda_graphs.empty() && time_now - last_graph_use_time >= 10'000'000) {
+            next_graph_eviction_time = time_now + 10'000'000;
+        }
+        last_graph_use_time = time_now;
 
-        // sweep every 5s, evicting cuda graphs unused for >=10s
-        if (time_now - last_graph_eviction_sweep >= 5'000'000) {
-            last_graph_eviction_sweep = time_now;
+        // Sweep every 5s during active use, evicting graphs unused for at least 10s.
+        if (time_now >= next_graph_eviction_time) {
+            next_graph_eviction_time = time_now + 5'000'000;
             for (auto it = cuda_graphs.begin(); it != cuda_graphs.end(); ) {
-                if (time_now - it->second->last_used_time >= 10'000'000) {
+                if (it->first != first_node_ptr && time_now - it->second->last_used_time >= 10'000'000) {
                     it = cuda_graphs.erase(it);
                 } else {
                     ++it;

ggml/src/ggml-cuda/fattn-common.cuh

@@ -25,6 +25,7 @@ typedef void (* fattn_kernel_t)(
         const char * __restrict__ mask,
         const char * __restrict__ sinks,
         const int  * __restrict__ KV_max,
+        const int  * __restrict__ KV_rows,
         float      * __restrict__ dst,
         float2     * __restrict__ dst_meta,
         const float scale,
@@ -991,8 +992,13 @@ void launch_fattn(
 
     const bool V_is_K_view = V->view_src && (V->view_src == K || (V->view_src == K->view_src && V->view_offs == K->view_offs));
 
-    const ggml_tensor * mask  = dst->src[3];
-    const ggml_tensor * sinks = dst->src[4];
+    const ggml_tensor * mask    = dst->src[3];
+    const ggml_tensor * sinks   = dst->src[4];
+    const ggml_tensor * kv_rows = dst->src[5];
+
+    // with kv_rows the kernel iterates over the mask columns and reads K/V rows through kv_rows
+    const int64_t n_kv_cols = kv_rows ? kv_rows->ne[0] : K->ne[1];
+    GGML_ASSERT(!kv_rows || !use_sparse);
 
     ggml_tensor * KQV = dst;
 
@@ -1113,8 +1119,8 @@ void launch_fattn(
 
     // Optional optimization where the mask is scanned to determine whether part of the calculation can be skipped.
     // Only worth the overhead if there is at lease one FATTN_KQ_STRIDE x FATTN_KQ_STRIDE square to be skipped or
-    //     multiple sequences of possibly different lengths.
-    const bool scan_mask = !use_sparse && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1);
+    //     multiple sequences of possibly different lengths, or kv_rows lists that end with padding.
+    const bool scan_mask = !use_sparse && mask && n_kv_cols % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1 || kv_rows);
     if (scan_mask) {
         const int64_t s31 = mask->nb[1] / sizeof(half2);
         const int64_t s33 = mask->nb[3] / sizeof(half2);
@@ -1123,7 +1129,7 @@ void launch_fattn(
         const dim3 block_dim_KV_max(FATTN_KQ_STRIDE/2, 1, 1);
 
         const int ne_KV_max = blocks_num_KV_max.x*blocks_num_KV_max.y;
-        const int iter_k = K->ne[1] / FATTN_KQ_STRIDE;
+        const int iter_k = n_kv_cols / FATTN_KQ_STRIDE;
 
         KV_max.alloc(ne_KV_max);
         ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks_num_KV_max, block_dim_KV_max, 0, main_stream);
@@ -1132,13 +1138,19 @@ void launch_fattn(
         CUDA_CHECK(cudaGetLastError());
     }
 
+    // K/V are shared by all slices when read through kv_rows
+    if (kv_rows) {
+        nb13 = 0;
+        nb23 = 0;
+    }
+
     const dim3 block_dim(warp_size, nwarps, 1);
     int max_blocks_per_sm = 1; // Max. number of active blocks limited by occupancy.
     CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&max_blocks_per_sm, fattn_kernel, block_dim.x * block_dim.y * block_dim.z, nbytes_shared));
     GGML_ASSERT(max_blocks_per_sm > 0);
     int parallel_blocks = max_blocks_per_sm;
 
-    const int64_t n_kv = use_sparse ? n_kv_max : K->ne[1];
+    const int64_t n_kv = use_sparse ? n_kv_max : n_kv_cols;
     const int ntiles_KV = (n_kv + nbatch_fa - 1) / nbatch_fa; // Max. number of parallel blocks limited by KV cache length.
 
     dim3 blocks_num;
@@ -1255,6 +1267,7 @@ void launch_fattn(
         mask ? ((const char *) mask->data) : nullptr,
         sinks ? ((const char *) sinks->data) : nullptr,
         KV_max.ptr,
+        kv_rows ? (const int *) kv_rows->data : nullptr,
         !stream_k && parallel_blocks > 1 ? dst_tmp.ptr : (float *) KQV->data, dst_tmp_meta.ptr,
         scale, max_bias, m0, m1, n_head_log2, logit_softcap,
         Q->ne[0], ne01,     Q->ne[2], Q->ne[3], Q->nb[1], Q->nb[2], Q->nb[3],

ggml/src/ggml-cuda/fattn-mma-f16.cuh

@@ -420,6 +420,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile(
                     // padded slots gather row 0, the -inf mask removes their contribution
                     const int32_t index = i < i_sup ? indices[k_VKQ_0 + i] : 0;
                     i_KV = index >= 0 ? index : 0;
+                } else if (indices) {
+                    // K/V rows given per mask column, negative rows are padding
+                    i_KV = max(indices[k_VKQ_0 + i], 0);
                 } else {
                     i_KV = k_VKQ_0 + i;
                 }
@@ -467,8 +470,11 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_tile(
                     if constexpr (use_sparse) {
                         const int32_t index = i < i_sup ? indices[k_VKQ_0 + i] : -1;
                         src = index >= 0 ? KV + int64_t(index)*stride_KV + k*h2_per_chunk : zero;
+                    } else if (!oob_check || i < i_sup) {
+                        const int64_t i_KV = indices ? max(indices[k_VKQ_0 + i], 0) : k_VKQ_0 + i;
+                        src = KV + i_KV*stride_KV + k*h2_per_chunk;
                     } else {
-                        src = !oob_check || i < i_sup ? KV + int64_t(k_VKQ_0 + i)*stride_KV + k*h2_per_chunk : zero;
+                        src = zero;
                     }
                     ggml_cuda_memcpy_1<16>(swizzle<stride_tile>(tile_KV, i*stride_tile + k*h2_per_chunk, i), src);
                 }
@@ -634,7 +640,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
         cp_async_wait_all();
         __syncthreads();
         flash_attn_ext_f16_load_tile<stride_tile_V, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
-            (V_h2, tile_V, nbatch_V2, stride_V, k_VKQ_0, k_VKQ_sup, nullptr);
+            (V_h2, tile_V, nbatch_V2, stride_V, k_VKQ_0, k_VKQ_sup, indices);
     } else {
         // the sparse mask values are gathered per element, always load them synchronously
         constexpr bool use_cp_async = nstages == 1 && !use_sparse;
@@ -991,7 +997,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter(
                     (mask_h, tile_mask, stride_mask, k_VKQ_0 + nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, nullptr);
             }
             flash_attn_ext_f16_load_tile<stride_tile_K, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
-                (K_h2, tile_K, nbatch_K2, stride_K, k_VKQ_0 + nbatch_fa, k_VKQ_sup, nullptr);
+                (K_h2, tile_K, nbatch_K2, stride_K, k_VKQ_0 + nbatch_fa, k_VKQ_sup, indices);
         }
     }
 
@@ -1255,7 +1261,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
     T_B_KQ    Q_B[(Q_in_reg ? DKQ/(2*T_B_KQ::J) : 1)];
 #if defined(TURING_MMA_AVAILABLE)
     T_C_VKQ VKQ_C[cols_per_warp == 8 ? DV/T_C_VKQ::I : DV/(2*T_C_VKQ::J)];
-#elif defined(AMD_WMMA_AVAILABLE) && defined(RDNA3)
+#elif defined(AMD_WMMA_256B_AVAILABLE)
     T_C_VKQ VKQ_C[DV % 32 != 0       ? DV/T_C_VKQ::J : DV/(2*T_C_VKQ::J)];
 #elif defined(AMD_MFMA_AVAILABLE)
     T_C_VKQ VKQ_C[                                     DV/T_C_VKQ::J];
@@ -1343,7 +1349,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile(
                 (mask_h, tile_mask, stride_mask, kb0*nbatch_fa, k_VKQ_sup, jt*ncols1, ne01, nullptr);
         }
         flash_attn_ext_f16_load_tile<stride_tile_K, nwarps, nbatch_fa, use_cp_async, oob_check, use_sparse>
-            (K_h2, tile_K, nbatch_K2, stride_K, kb0*nbatch_fa, k_VKQ_sup, nullptr);
+            (K_h2, tile_K, nbatch_K2, stride_K, kb0*nbatch_fa, k_VKQ_sup, indices);
     }
 
     // kb0_start is always < kb0_stop so the last iter can be executed unconditionally.
@@ -1802,6 +1808,7 @@ static __global__ void flash_attn_ext_f16(
         const char * mask_ptr,
         const char * sinks_ptr,
         const int  * KV_max_ptr,
+        const int  * KV_rows_ptr,
         float      * dst_ptr,
         float2     * dst_meta_ptr,
         const float scale,
@@ -1866,7 +1873,7 @@ static __global__ void flash_attn_ext_f16(
 #endif // defined(AMD_WMMA_AVAILABLE)
 
 #if defined(AMD_MFMA_AVAILABLE)
-    if (ncols1*ncols2 < 16 || (DKQ > 256 && ncols1*ncols2 < 32)) {
+    if (ncols1*ncols2 < 16 || (DKQ > 256 && ncols1*ncols2 < 32 && !use_sparse)) {
         NO_DEVICE_CODE;
         return;
     }
@@ -1924,7 +1931,8 @@ static __global__ void flash_attn_ext_f16(
 
         const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV);
         const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr;
-        const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*iter_j + jt)*ne11 : nullptr;
+        const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*iter_j + jt)*ne11 :
+                                  KV_rows_ptr ? KV_rows_ptr + int64_t(sequence)*ne11 : nullptr;
 
         const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f;
 
@@ -1973,7 +1981,8 @@ static __global__ void flash_attn_ext_f16(
 
     const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV);
     const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr;
-    const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*iter_j + jt)*ne11 : nullptr;
+    const int32_t * indices = use_sparse ? sparse_indices + (int64_t(sequence % ne33)*iter_j + jt)*ne11 :
+                              KV_rows_ptr ? KV_rows_ptr + int64_t(sequence)*ne11 : nullptr;
 
     const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f;
 
@@ -1989,7 +1998,7 @@ static __global__ void flash_attn_ext_f16(
         (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap,
          ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop);
 #else
-    GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale,
+    GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, KV_rows_ptr, dst_ptr, dst_meta_ptr, scale,
         max_bias, m0, m1, n_head_log2, logit_softcap,
         ne00, ne01, ne02, ne03,
               nb01, nb02, nb03,
@@ -2054,7 +2063,7 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml
     bool use_sparse = false;
     if (logit_softcap == 0.0f) {
         constexpr bool use_logit_softcap = false;
-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#if !defined(GGML_USE_MUSA)
         if constexpr (ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, ncols1, ncols2)) {
             if (ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(cc, dst, ncols1, ncols2)) {
                 constexpr bool use_sparse_kernel = true;
@@ -2077,7 +2086,7 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml
                 }
             }
         } else
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#endif // !defined(GGML_USE_MUSA)
         {
             constexpr bool use_sparse_kernel = false;
             fattn_kernel = flash_attn_ext_f16<DKQ, DV, ncols1, ncols2, use_logit_softcap, V_is_K_view, use_sparse_kernel>;

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

@@ -797,6 +797,7 @@ static __global__ void flash_attn_tile(
         const char * mask_ptr,
         const char * sinks_ptr,
         const int  * KV_max_ptr,
+        const int  * KV_rows_ptr,
         float      * dst_ptr,
         float2     * dst_meta_ptr,
         const float scale,
@@ -821,6 +822,7 @@ static __global__ void flash_attn_tile(
     const int  * GGML_CUDA_RESTRICT KV_max   = KV_max_ptr;
     float      * GGML_CUDA_RESTRICT dst      = dst_ptr;
     float2     * GGML_CUDA_RESTRICT dst_meta = dst_meta_ptr;
+    GGML_UNUSED(KV_rows_ptr);
 
     // Skip unused kernel variants for faster compilation:
 
@@ -1132,7 +1134,7 @@ static __global__ void flash_attn_tile(
         }
     }
 #else
-    GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale,
+    GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, KV_rows_ptr, dst_ptr, dst_meta_ptr, scale,
         max_bias, m0, m1, n_head_log2, logit_softcap,
         ne00, ne01, ne02, ne03,
               nb01, nb02, nb03,

ggml/src/ggml-cuda/fattn-vec.cuh

@@ -25,6 +25,7 @@ static __global__ void flash_attn_ext_vec(
         const char * mask_ptr,
         const char * sinks_ptr,
         const int  * KV_max_ptr,
+        const int  * KV_rows_ptr,
         float      * dst_ptr,
         float2     * dst_meta_ptr,
         const float scale,
@@ -50,6 +51,7 @@ static __global__ void flash_attn_ext_vec(
     const int  * GGML_CUDA_RESTRICT KV_max   = KV_max_ptr;
     float      * GGML_CUDA_RESTRICT dst      = dst_ptr;
     float2     * GGML_CUDA_RESTRICT dst_meta = dst_meta_ptr;
+    GGML_UNUSED(KV_rows_ptr);
 
     // Skip unused kernel variants for faster compilation:
     if (use_logit_softcap && !(D == 128 || D == 256)) {
@@ -514,7 +516,7 @@ static __global__ void flash_attn_ext_vec(
         dst_meta[((sequence*int(ne01.z) + ic0 + tid)*ne02 + head)*gridDim.y + blockIdx.y] = make_float2(KQ_max[tid], KQ_sum[tid]);
     }
 #else
-    GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale,
+    GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, KV_rows_ptr, dst_ptr, dst_meta_ptr, scale,
         max_bias, m0, m1, n_head_log2, logit_softcap,
         ne00, ne01, ne02, ne03,
               nb01, nb02, nb03,

ggml/src/ggml-cuda/fattn.cu

@@ -5,7 +5,8 @@
 #include "fattn-vec.cuh"
 #include "fattn.cuh"
 
-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#if !defined(GGML_USE_MUSA)
+
 // one list per group of ncols1 queries: a column is selected if any query of the group can see it
 template <int ncols1, bool oob>
 __launch_bounds__(256, 1)
@@ -14,10 +15,20 @@ static __global__ void flash_attn_mask_to_sparse_indices(
         const int n_kv_max, const int64_t s31, const int64_t s33) {
     ggml_cuda_pdl_sync();
 
+    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
     constexpr int values_per_lane = 8;
+    using lane_mask_t = std::conditional_t<warp_size == 64, uint64_t, uint32_t>;
+    constexpr uint64_t ballot_mask = ~(lane_mask_t)0;
+    const auto popc = [](const lane_mask_t mask) {
+        if constexpr (warp_size == 64) {
+            return __popcll(mask);
+        } else {
+            return __popc(mask);
+        }
+    };
     const int tid      = threadIdx.x;
-    const int warp     = tid / WARP_SIZE;
-    const int lane     = tid % WARP_SIZE;
+    const int warp     = tid / warp_size;
+    const int lane     = tid % warp_size;
     const int sequence = blockIdx.y;
     const int group    = blockIdx.x;
 
@@ -27,7 +38,7 @@ static __global__ void flash_attn_mask_to_sparse_indices(
     const half * mask = mask_ptr + sequence*s33 + q0*s31;
     int32_t * indices = indices_ptr + (int64_t(sequence)*gridDim.x + group)*n_kv_max;
 
-    __shared__ int warp_offsets[256/WARP_SIZE];
+    __shared__ int warp_offsets[256/warp_size];
     __shared__ int row_count;
     __shared__ int chunk_count;
 
@@ -37,11 +48,11 @@ static __global__ void flash_attn_mask_to_sparse_indices(
     __syncthreads();
 
     for (int i0 = 0; i0 < ne30; i0 += blockDim.x*values_per_lane) {
-        uint32_t selected_warp[values_per_lane];
+        lane_mask_t selected_warp[values_per_lane];
         int warp_count = 0;
 #pragma unroll
         for (int item = 0; item < values_per_lane; ++item) {
-            const int i = i0 + (warp*values_per_lane + item)*WARP_SIZE + lane;
+            const int i = i0 + (warp*values_per_lane + item)*warp_size + lane;
             bool selected = false;
             if (i < ne30) {
 #pragma unroll
@@ -49,8 +60,8 @@ static __global__ void flash_attn_mask_to_sparse_indices(
                     selected |= (!oob || q < q1 - q0) && isfinite(__half2float(mask[q*s31 + i]));
                 }
             }
-            selected_warp[item] = __ballot_sync(0xFFFFFFFF, selected);
-            warp_count += __popc(selected_warp[item]);
+            selected_warp[item] = (lane_mask_t)__ballot_sync(ballot_mask, (int) selected);
+            warp_count += popc(selected_warp[item]);
         }
 
         if (lane == 0) {
@@ -61,7 +72,7 @@ static __global__ void flash_attn_mask_to_sparse_indices(
         if (tid == 0) {
             int offset = 0;
 #pragma unroll
-            for (int iw = 0; iw < 256/WARP_SIZE; ++iw) {
+            for (int iw = 0; iw < 256/warp_size; ++iw) {
                 const int count = warp_offsets[iw];
                 warp_offsets[iw] = offset;
                 offset += count;
@@ -70,16 +81,16 @@ static __global__ void flash_attn_mask_to_sparse_indices(
         }
         __syncthreads();
 
-        const uint32_t lane_mask = lane == 0 ? 0 : (1u << lane) - 1;
+        const lane_mask_t lane_mask = lane == 0 ? 0 : ((lane_mask_t)1 << lane) - 1;
         int warp_item_offset = 0;
 #pragma unroll
         for (int item = 0; item < values_per_lane; ++item) {
-            const int i = i0 + (warp*values_per_lane + item)*WARP_SIZE + lane;
-            const int dst = row_count + warp_offsets[warp] + warp_item_offset + __popc(selected_warp[item] & lane_mask);
-            if ((selected_warp[item] & (uint32_t(1) << lane)) && dst < n_kv_max) {
+            const int i = i0 + (warp*values_per_lane + item)*warp_size + lane;
+            const int dst = row_count + warp_offsets[warp] + warp_item_offset + popc(selected_warp[item] & lane_mask);
+            if ((selected_warp[item] & ((lane_mask_t)1 << lane)) && dst < n_kv_max) {
                 indices[dst] = i;
             }
-            warp_item_offset += __popc(selected_warp[item]);
+            warp_item_offset += popc(selected_warp[item]);
         }
         __syncthreads();
 
@@ -101,13 +112,13 @@ static __global__ void flash_attn_mask_to_sparse_indices(
     // the dependent grid reads indices, signal once the row is complete
     ggml_cuda_pdl_lc();
 }
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#endif // !defined(GGML_USE_MUSA)
 
 void ggml_cuda_flash_attn_ext_compact_mask(
         const ggml_tensor * mask, int32_t * indices, int32_t * counts, int32_t n_queries, int32_t ncols1, int32_t n_kv_max, cudaStream_t stream) {
-#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
+#ifdef GGML_USE_MUSA
     GGML_UNUSED_VARS(mask, indices, counts, n_queries, ncols1, n_kv_max, stream);
-    GGML_ABORT("sparse flash attention is only supported on NVIDIA CUDA");
+    GGML_ABORT("sparse flash attention is not supported on MUSA");
 #else
     const int64_t s31 = mask->nb[1] / sizeof(half);
     const int64_t s33 = mask->nb[3] / sizeof(half);
@@ -122,16 +133,17 @@ void ggml_cuda_flash_attn_ext_compact_mask(
     ggml_cuda_kernel_launch(kernel, launch_params,
         (const half *) mask->data, indices, counts, int(mask->ne[0]), n_queries, n_kv_max, s31, s33);
     CUDA_CHECK(cudaGetLastError());
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#endif // !defined(GGML_USE_MUSA)
 }
 
 bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(const int cc, const ggml_tensor * dst, const int ncols1, const int ncols2) {
-#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
+#ifdef GGML_USE_MUSA
     GGML_UNUSED_VARS(cc, dst, ncols1, ncols2);
     return false;
 #else
     const ggml_tensor * Q    = dst->src[0];
     const ggml_tensor * K    = dst->src[1];
+    const ggml_tensor * V    = dst->src[2];
     const ggml_tensor * mask = dst->src[3];
 
     float max_bias = 0.0f;
@@ -139,16 +151,21 @@ bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(const int cc, const ggml_
     memcpy(&max_bias,      (const float *) dst->op_params + 1, sizeof(float));
     memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float));
 
+    // On CDNA only the DKQ=576 kernel variant with ncols2=16 is compiled with sparse support.
+    if (amd_mfma_available(cc) && !ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(int(Q->ne[0]), int(V->ne[0]), 1, 16)) {
+        return false;
+    }
+
     const int32_t n_kv_max = ggml_get_op_params_i32(dst, 4);
 
     // the dense kernel handles up to 64/ncols2 queries per K/V pass, the single-query gather has to beat that
     const int64_t n_gather = (ncols1 == 1 ? std::min<int64_t>(Q->ne[1], 64/ncols2) : ncols1) * (int64_t) n_kv_max;
 
-    return GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) &&
+    return ((GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc)) || amd_mfma_available(cc)) && dst->src[5] == nullptr &&
         mask != nullptr && n_kv_max > 0 && max_bias == 0.0f && logit_softcap == 0.0f &&
         mask->ne[0] == K->ne[1] && mask->ne[1] >= Q->ne[1] && mask->ne[2] == 1 &&
-        K->ne[1] >= std::max<int64_t>(4096, 2*n_gather);
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+        K->ne[1] >= std::max<int64_t>(4096, GGML_CUDA_CC_IS_NVIDIA(cc) ? 2*n_gather : 2LL*n_kv_max);
+#endif // !defined(GGML_USE_MUSA)
 }
 
 template <int DKQ, int DV, int ncols2>
@@ -156,7 +173,7 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con
     const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
     const ggml_tensor * Q = dst->src[0];
 
-#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#if !defined(GGML_USE_MUSA)
     if constexpr (ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, 1, ncols2)) {
         // a sparse variant at the full tile width gathers the union of its queries once, prefer it for large batches
         constexpr bool has_wide_sparse = ggml_cuda_flash_attn_ext_mma_f16_may_use_sparse(DKQ, DV, 64/ncols2, ncols2);
@@ -165,7 +182,7 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con
             return;
         }
     }
-#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#endif // !defined(GGML_USE_MUSA)
 
     if constexpr (ncols2 <= 8) {
         if (turing_mma_available(cc) && Q->ne[1] <= 8/ncols2) {
@@ -199,12 +216,15 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols2(ggml_backend_cuda_con
     const ggml_tensor * V    = dst->src[2];
     const ggml_tensor * mask = dst->src[3];
 
+    // with kv_rows the kernel iterates over the mask columns
+    const int64_t n_kv = dst->src[5] ? dst->src[5]->ne[0] : K->ne[1];
+
     float max_bias = 0.0f;
     memcpy(&max_bias, (const float *) KQV->op_params + 1, sizeof(float));
 
     // Edge cases like no mask, ALiBi, unpadded K/V, or misaligned addresses for large data transfers
     //     are put into the template specialization without GQA optimizations.
-    bool use_gqa_opt = mask && max_bias == 0.0f && K->ne[1] % FATTN_KQ_STRIDE == 0;
+    bool use_gqa_opt = mask && max_bias == 0.0f && n_kv % FATTN_KQ_STRIDE == 0;
     for (const ggml_tensor * t : {Q, K, V, mask}) {
         if (t == nullptr || ggml_is_quantized(t->type)) {
             continue;
@@ -365,7 +385,9 @@ 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 (amd_mfma_available(cc) && ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(cc, dst, 1, 16)) {
+                ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1<576, 512, 16>(ctx, dst);
+            } else if (gqa_ratio == 20 && GGML_CUDA_CC_IS_NVIDIA(cc)) { // 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);
@@ -548,6 +570,10 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
     const ggml_tensor * K     = dst->src[1];
     const ggml_tensor * V     = dst->src[2];
     const ggml_tensor * mask  = dst->src[3];
+    const ggml_tensor * kv_rows = dst->src[5];
+
+    // with kv_rows the kernel iterates over the mask columns
+    const int64_t n_kv = kv_rows ? kv_rows->ne[0] : K->ne[1];
 
     const int gqa_ratio = Q->ne[2] / K->ne[2];
     GGML_ASSERT(Q->ne[2] % K->ne[2] == 0);
@@ -557,7 +583,7 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
 
     // The effective batch size for the kernel can be increased by gqa_ratio.
     // The kernel versions without this optimization are also used for ALiBi, if there is no mask, or if the KV cache is not padded,
-    bool gqa_opt_applies = gqa_ratio >= 2 && mask && max_bias == 0.0f && K->ne[1] % FATTN_KQ_STRIDE == 0;
+    bool gqa_opt_applies = gqa_ratio >= 2 && mask && max_bias == 0.0f && n_kv % FATTN_KQ_STRIDE == 0;
     for (const ggml_tensor * t : {Q, K, V, mask}) {
         if (t == nullptr || ggml_is_quantized(t->type)) {
             continue;
@@ -629,6 +655,12 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
         return BEST_FATTN_KERNEL_NONE;
     }
 
+    // only the MMA kernel reads K/V through kv_rows
+    if (kv_rows) {
+        return GGML_CUDA_CC_IS_NVIDIA(cc) && turing_mma_available(cc) && mask && Q->ne[0] != 40 && Q->ne[0] != 72 ?
+            BEST_FATTN_KERNEL_MMA_F16 : BEST_FATTN_KERNEL_NONE;
+    }
+
     // For small batch sizes the vector kernel may be preferable over the kernels optimized for large batch sizes:
     // 192 satisfies % 64 == 0 but has no vec instance (DKQ != DV); force it onto the MMA path.
     const bool can_use_vector_kernel = Q->ne[0] <= 256 && Q->ne[0] % 64 == 0 && Q->ne[0] != 192 && K->ne[1] % FATTN_KQ_STRIDE == 0;
@@ -679,6 +711,11 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
         return BEST_FATTN_KERNEL_MMA_F16;
     }
 
+    // On CDNA, the sparse MMA kernel is used for DKQ=576 regardless of the head size heuristics.
+    if (amd_mfma_available(cc) && ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(cc, dst, 1, 16)) {
+        return BEST_FATTN_KERNEL_MMA_F16;
+    }
+
     // AMD MFMA needs a certain minimum batch size to outscale the tile kernel for large head sizes.
     if (amd_mfma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72) {
         if ((Q->ne[0] <= 64 && Q->ne[1] * gqa_ratio_eff > 8)) {

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

@@ -9,6 +9,7 @@
 #include "ggml-cuda/arange.cuh"
 #include "ggml-cuda/argmax.cuh"
 #include "ggml-cuda/argsort.cuh"
+#include "ggml-cuda/batch-get.cuh"
 #include "ggml-cuda/binbcast.cuh"
 #include "ggml-cuda/clamp.cuh"
 #include "ggml-cuda/col2im-1d.cuh"
@@ -861,6 +862,51 @@ static void ggml_backend_cuda_buffer_get_tensor_2d(ggml_backend_buffer_t buffer,
     CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
 }
 
+static void ggml_backend_cuda_buffer_set_tensor_n(ggml_backend_buffer_t buffer, struct ggml_tensor * const * tensors, const void * const * data, const size_t * offsets, const size_t * sizes, size_t n_tensors) {
+    ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context;
+
+    ggml_cuda_set_device(ctx->device);
+
+    for (size_t i = 0; i < n_tensors; i++) {
+        if (sizes[i] == 0) {
+            continue;
+        }
+        CUDA_CHECK(cudaMemcpyAsync(
+            (char *) tensors[i]->data + offsets[i],
+            data[i],
+            sizes[i],
+            cudaMemcpyHostToDevice,
+            cudaStreamPerThread));
+    }
+
+    CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
+}
+
+static void ggml_backend_cuda_buffer_get_tensor_n(ggml_backend_buffer_t buffer, const struct ggml_tensor * const * tensors, void * const * data, const size_t * offsets, const size_t * sizes, size_t n_tensors) {
+    ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context;
+
+    ggml_cuda_set_device(ctx->device);
+
+    if (!ggml_cuda_pinned_memory_disabled() &&
+            ggml_cuda_get_tensor_n(ctx->device, tensors, data, offsets, sizes, n_tensors)) {
+        return;
+    }
+
+    for (size_t i = 0; i < n_tensors; i++) {
+        if (sizes[i] == 0) {
+            continue;
+        }
+        CUDA_CHECK(cudaMemcpyAsync(
+            data[i],
+            (const char *) tensors[i]->data + offsets[i],
+            sizes[i],
+            cudaMemcpyDeviceToHost,
+            cudaStreamPerThread));
+    }
+
+    CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
+}
+
 static bool ggml_backend_cuda_buffer_cpy_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * src, ggml_tensor * dst) {
     if (ggml_backend_buffer_is_cuda(src->buffer)) {
         ggml_backend_cuda_buffer_context * src_ctx = (ggml_backend_cuda_buffer_context *)src->buffer->context;
@@ -886,6 +932,54 @@ static bool ggml_backend_cuda_buffer_cpy_tensor(ggml_backend_buffer_t buffer, co
     GGML_UNUSED(buffer);
 }
 
+static bool ggml_backend_cuda_buffer_cpy_tensor_n(ggml_backend_buffer_t buffer, const struct ggml_tensor * const * srcs, struct ggml_tensor * const * dsts, size_t n_tensors) {
+    ggml_backend_cuda_buffer_context * dst_ctx = (ggml_backend_cuda_buffer_context *)buffer->context;
+
+    for (size_t i = 0; i < n_tensors; i++) {
+        ggml_backend_buffer_t src_buf = srcs[i]->view_src ? srcs[i]->view_src->buffer : srcs[i]->buffer;
+        if (!ggml_backend_buffer_is_cuda(src_buf)) {
+            return false;
+        }
+    }
+
+#ifdef GGML_CUDA_NO_PEER_COPY
+    for (size_t i = 0; i < n_tensors; i++) {
+        ggml_backend_buffer_t src_buf = srcs[i]->view_src ? srcs[i]->view_src->buffer : srcs[i]->buffer;
+        ggml_backend_buffer_t dst_buf = dsts[i]->view_src ? dsts[i]->view_src->buffer : dsts[i]->buffer;
+        ggml_backend_cuda_buffer_context * src_ctx = (ggml_backend_cuda_buffer_context *)src_buf->context;
+        ggml_backend_cuda_buffer_context * dst_ctx_i = (ggml_backend_cuda_buffer_context *)dst_buf->context;
+        const int src_physical = ggml_cuda_get_physical_device(src_ctx->device);
+        const int dst_physical = ggml_cuda_get_physical_device(dst_ctx_i->device);
+        if (src_physical != dst_physical) {
+            return false;
+        }
+    }
+#endif
+
+    ggml_cuda_set_device(dst_ctx->device);
+
+    for (size_t i = 0; i < n_tensors; i++) {
+        ggml_backend_buffer_t src_buf = srcs[i]->view_src ? srcs[i]->view_src->buffer : srcs[i]->buffer;
+        ggml_backend_buffer_t dst_buf = dsts[i]->view_src ? dsts[i]->view_src->buffer : dsts[i]->buffer;
+        ggml_backend_cuda_buffer_context * src_ctx = (ggml_backend_cuda_buffer_context *)src_buf->context;
+        ggml_backend_cuda_buffer_context * dst_ctx_i = (ggml_backend_cuda_buffer_context *)dst_buf->context;
+        const int src_physical = ggml_cuda_get_physical_device(src_ctx->device);
+        const int dst_physical = ggml_cuda_get_physical_device(dst_ctx_i->device);
+        if (src_physical == dst_physical) {
+            CUDA_CHECK(cudaMemcpyAsync(dsts[i]->data, srcs[i]->data, ggml_nbytes(srcs[i]), cudaMemcpyDeviceToDevice, cudaStreamPerThread));
+        } else {
+#ifdef GGML_CUDA_NO_PEER_COPY
+            GGML_ABORT("peer copy not supported");
+#else
+            CUDA_CHECK(cudaMemcpyPeerAsync(dsts[i]->data, dst_physical, srcs[i]->data, src_physical, ggml_nbytes(srcs[i]), cudaStreamPerThread));
+#endif
+        }
+    }
+
+    CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
+    return true;
+}
+
 static void ggml_backend_cuda_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) {
     ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context;
 
@@ -906,6 +1000,9 @@ static const ggml_backend_buffer_i ggml_backend_cuda_buffer_interface = {
     /* .cpy_tensor      = */ ggml_backend_cuda_buffer_cpy_tensor,
     /* .clear           = */ ggml_backend_cuda_buffer_clear,
     /* .reset           = */ NULL,
+    /* .set_tensor_n    = */ ggml_backend_cuda_buffer_set_tensor_n,
+    /* .get_tensor_n    = */ ggml_backend_cuda_buffer_get_tensor_n,
+    /* .cpy_tensor_n    = */ ggml_backend_cuda_buffer_cpy_tensor_n,
 };
 
 // cuda buffer type
@@ -1026,6 +1123,18 @@ struct ggml_backend_cuda_comm_context {
 
     ggml_cuda_ar_pipeline *     ar_pipeline = nullptr;
 
+    // Number of ranks in the communicator when each rank is a separate process. Zero unless
+    // comm_init_rank created this context.
+    int                         world_size = 0;
+
+    // Number of ranks taking part in the reduction. For a communicator from comm_init_rank this is
+    // world_size. Otherwise it is the number of backends this process holds. The NCCL reduction
+    // uses it to choose the tensor size below which it reduces in F32 rather than converting to
+    // BF16 first.
+    size_t n_ranks() const {
+        return world_size > 0 ? (size_t) world_size : backends.size();
+    }
+
 #ifdef GGML_USE_NCCL
     std::vector<ncclComm_t>     comms;
 #endif // GGML_USE_NCCL
@@ -1053,6 +1162,7 @@ static bool ggml_backend_cuda_comm_allreduce_nccl(
     }
 
     const size_t n_backends = comm_ctx->backends.size();
+    const size_t n_ranks    = comm_ctx->n_ranks();
 
     for (size_t i = 0; i < n_backends; ++i) {
         GGML_ASSERT(tensors[i] != nullptr);
@@ -1062,7 +1172,7 @@ static bool ggml_backend_cuda_comm_allreduce_nccl(
 
     // For small tensors, simply reduce them as FP32.
     // The following heuristic for how "small" a tensor should be is based on RTX 4090s connected via 16x PCIe 4.0.
-    if ((n_backends <= 2 && ne < 32768) || (n_backends == 3 && ne < 131072) || (n_backends >= 4 && ne < 262144)) {
+    if ((n_ranks <= 2 && ne < 32768) || (n_ranks == 3 && ne < 131072) || (n_ranks >= 4 && ne < 262144)) {
         for (size_t i = 0; i < n_backends; ++i) {
             if ((tensors[i]->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) {
                 ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) comm_ctx->backends[i]->context;
@@ -1293,6 +1403,90 @@ static void * ggml_backend_cuda_comm_init(ggml_backend_t * backends, size_t n_ba
     return ret;
 }
 
+#ifdef GGML_USE_NCCL
+// Builds the communicator with ncclCommInitRank, for the case where each rank is its own
+// process. Blocks until all world_size ranks have joined.
+static void * ggml_backend_cuda_comm_init_rank(ggml_backend_t backend, int rank, int world_size,
+                                               ggml_backend_comm_bootstrap_callback bootstrap_callback, void * bootstrap_callback_data) {
+    if (world_size < 1 || rank < 0 || rank >= world_size) {
+        GGML_LOG_ERROR("%s: invalid rank %d for world size %d\n", __func__, rank, world_size);
+        return nullptr;
+    }
+
+    // Every rank sends whether its local checks passed, and rank 0 also sends the NCCL unique id.
+    // So all ranks either join together or fail together, and no rank waits in ncclCommInitRank alone.
+    struct bootstrap_msg {
+        uint8_t      ok;
+        ncclUniqueId id;
+    };
+    bootstrap_msg local;
+    memset(&local, 0, sizeof(local));
+    local.ok = 1;
+
+    if (!ggml_backend_is_cuda(backend)) {
+        GGML_LOG_ERROR("%s: backend is not a CUDA backend\n", __func__);
+        local.ok = 0;
+    } else {
+        // NCCL requires one distinct physical GPU per rank, see ggml_backend_cuda_comm_init_nccl.
+        const ggml_cuda_device_info & info = ggml_cuda_info();
+        if (info.device_count > info.physical_device_count) {
+            GGML_LOG_ERROR("%s: CUDA virtual devices are in use, which NCCL does not support\n", __func__);
+            local.ok = 0;
+        }
+    }
+    if (local.ok && rank == 0) {
+        const ncclResult_t rc = ncclGetUniqueId(&local.id);
+        if (rc != ncclSuccess) {
+            GGML_LOG_ERROR("%s: ncclGetUniqueId failed (%s)\n", __func__, ncclGetErrorString(rc));
+            local.ok = 0;
+        }
+    }
+
+    std::vector<bootstrap_msg> all(world_size);
+    if (!bootstrap_callback(bootstrap_callback_data, &local, all.data(), sizeof(local))) {
+        GGML_LOG_ERROR("%s: bootstrap exchange failed on rank %d\n", __func__, rank);
+        return nullptr;
+    }
+    for (int i = 0; i < world_size; i++) {
+        if (!all[i].ok) {
+            GGML_LOG_ERROR("%s: rank %d failed its setup checks\n", __func__, i);
+            return nullptr;
+        }
+    }
+
+    ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context;
+
+    auto * ret = new ggml_backend_cuda_comm_context;
+    ret->backends.push_back(backend);
+    ret->dev_ids.push_back(cuda_ctx->device);
+    ret->world_size = world_size;
+    ret->comms.resize(1);
+    ret->try_allreduce = ggml_backend_cuda_comm_try_allreduce_nccl;
+
+    // On DGX Spark, NCCL picks 64 channels by default, which is slower and uses more memory than 8.
+    // NCCL_MAX_CTAS still overrides this.
+    ncclConfig_t config = NCCL_CONFIG_INITIALIZER;
+    const int cc = ggml_cuda_info().devices[cuda_ctx->device].cc;
+    if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_DGX_SPARK) {
+        config.maxCTAs = 8;
+    }
+
+    ggml_cuda_set_device(cuda_ctx->device);
+    const ncclResult_t rc = ncclCommInitRankConfig(&ret->comms[0], world_size, all[0].id, rank, &config);
+    if (rc != ncclSuccess) {
+        GGML_LOG_ERROR("%s: ncclCommInitRankConfig failed for rank %d/%d (%s)\n",
+                       __func__, rank, world_size, ncclGetErrorString(rc));
+        ret->comms.clear();
+        delete ret;
+        return nullptr;
+    }
+
+    GGML_LOG_INFO("%s: joined NCCL communicator as rank %d/%d on device %d\n",
+                  __func__, rank, world_size, cuda_ctx->device);
+    return ret;
+}
+#endif // GGML_USE_NCCL
+
 // Top-level dispatch -- calls the function pointer chosen by comm_init.
 // Returns false to let the meta-backend's butterfly run.
 static bool ggml_backend_cuda_comm_allreduce_tensor(void * comm_ctx_v, struct ggml_tensor ** tensors) {
@@ -1565,7 +1759,7 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const
     if (compute_type == GGML_TYPE_F16) {
         prefer_f32_output = cc == GGML_CUDA_CC_VOLTA || GGML_CUDA_CC_IS_RDNA4(cc) || GGML_CUDA_CC_IS_CDNA(cc);
     } else if (compute_type == GGML_TYPE_BF16) {
-        prefer_f32_output = !GGML_CUDA_CC_IS_RDNA3(cc) && !GGML_CUDA_CC_IS_CDNA(cc);
+        prefer_f32_output = !GGML_CUDA_CC_IS_RDNA3(cc) && !GGML_CUDA_CC_IS_RDNA4_M(cc) && !GGML_CUDA_CC_IS_CDNA(cc);
     }
 
     if (prefer_f32_output) {
@@ -6293,6 +6487,11 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con
     if (strcmp(name, "ggml_backend_comm_allreduce_tensor") == 0) {
         return (void *)ggml_backend_cuda_comm_allreduce_tensor;
     }
+#ifdef GGML_USE_NCCL
+    if (strcmp(name, "ggml_backend_comm_init_rank") == 0) {
+        return (void *)ggml_backend_cuda_comm_init_rank;
+    }
+#endif // GGML_USE_NCCL
     if (strcmp(name, "ggml_backend_register_host_buffer") == 0) {
         return (void *)ggml_backend_cuda_register_host_buffer;
     }

ggml/src/ggml-cuda/mma.cuh

@@ -76,11 +76,11 @@ namespace ggml_cuda_mma {
         // For the A/C matrices this means I major == row major, J major == column major.
         // For the B matrix this means I major == column major, J major == row major.
         // MIRRORED == Each data value is held exactly once per thread subgroup.
-        DATA_LAYOUT_I_MAJOR           =  0, // Always used for Turing, Ampere, Ada Lovelace, consumer Blackwell, matrix A&B for RDNA4 and CDNA.
-        DATA_LAYOUT_J_MAJOR           = 10, // Matrix C for CDNA and RDNA4, int and float matrix C for RDNA3.
+        DATA_LAYOUT_I_MAJOR           =  0, // Always used for Turing, Ampere, Ada Lovelace, consumer Blackwell, matrix A&B for WMMA128 and CDNA.
+        DATA_LAYOUT_J_MAJOR           = 10, // Matrix C for CDNA and WMMA128, int and float matrix C for WMMA256.
         DATA_LAYOUT_I_MAJOR_MIRRORED  = 20, // Volta, matrix A&B for RDNA3.
         DATA_LAYOUT_J_MAJOR_MIRRORED  = 30,
-        DATA_LAYOUT_I_MAJOR_SCRAMBLED = 40, // Scrambled matrix C for faster transposition (RDNA4/CDNA), convert to float to unscramble.
+        DATA_LAYOUT_I_MAJOR_SCRAMBLED = 40, // Scrambled matrix C for faster transposition (WMMA128/CDNA), convert to float to unscramble.
     };
     // Implemented mma combinations are:
     //   - (I_MAJOR, I_MAJOR)          -> I_MAJOR
@@ -88,11 +88,11 @@ namespace ggml_cuda_mma {
     //   - (I_MAJOR, J_MAJOR_MIRRORED) -> I_MAJOR
 
     static constexpr __device__ data_layout get_input_data_layout() {
-#if defined(RDNA3) || defined(VOLTA_MMA_AVAILABLE)
+#if defined(AMD_WMMA_256B_AVAILABLE) || defined(VOLTA_MMA_AVAILABLE)
         return DATA_LAYOUT_I_MAJOR_MIRRORED;
 #else
         return DATA_LAYOUT_I_MAJOR;
-#endif // defined(RDNA3) || defined(VOLTA_MMA_AVAILABLE)
+#endif // defined(AMD_WMMA_256B_AVAILABLE) || defined(VOLTA_MMA_AVAILABLE)
     }
 
     template <int I_, int J_, typename T, data_layout ds_=DATA_LAYOUT_I_MAJOR>
@@ -201,7 +201,7 @@ namespace ggml_cuda_mma {
 
         static __device__ __forceinline__ int get_j(const int l) {
             if constexpr (I == 16 && J == 16) {
-#if defined(RDNA3)
+#if defined(AMD_WMMA_256B_AVAILABLE)
                 if constexpr (std::is_same_v<T, float> || std::is_same_v<T, int>) {
                     // matrix C
                     return 2 * l + (threadIdx.x / 16);
@@ -210,11 +210,11 @@ namespace ggml_cuda_mma {
                     return l;
                 }
 #else
-                // matrix C is the transposed matrix A&B on RDNA4
+                // matrix C is the transposed matrix A&B on WMMA128
                 return ne * (threadIdx.x / 16) + l;
-#endif // defined(RDNA3)
+#endif // defined(AMD_WMMA_256B_AVAILABLE)
             } else if constexpr (I == 16 && J == 8) {
-                // mmq input for RDNA4
+                // mmq input for WMMA128
                 return ne * (threadIdx.x / 16) + l;
             } else if constexpr (I == 16 && J == 4) {
                 return ne * (threadIdx.x / 16) + l;
@@ -573,7 +573,7 @@ namespace ggml_cuda_mma {
         static constexpr int         I  = I_;
         static constexpr int         J  = J_;
         static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR_MIRRORED;
-#if defined(RDNA3)
+#if defined(AMD_WMMA_256B_AVAILABLE)
         static constexpr int         ne = tile<I_, J_, float, DATA_LAYOUT_I_MAJOR_MIRRORED>::ne;
 
         half2 x[ne] = {{0.0f, 0.0f}};
@@ -616,7 +616,7 @@ namespace ggml_cuda_mma {
                 return -1;
             }
         }
-#endif // defined(RDNA3)
+#endif // defined(AMD_WMMA_256B_AVAILABLE)
     };
 
     template <int I_, int J_>
@@ -694,7 +694,7 @@ namespace ggml_cuda_mma {
     };
 
     static __device__ __forceinline__ tile<16, 16, half2, DATA_LAYOUT_I_MAJOR> unscramble(const tile<16, 16, half2, DATA_LAYOUT_I_MAJOR_SCRAMBLED> & t) {
-#if defined(AMD_MFMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA4))
+#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_128B_AVAILABLE)
         tile<16, 16, half2, DATA_LAYOUT_I_MAJOR> ret;
 #pragma unroll
         for (int l0 = 0; l0 < t.ne/2; ++l0) {
@@ -705,7 +705,7 @@ namespace ggml_cuda_mma {
 #else
         NO_DEVICE_CODE;
         GGML_UNUSED(t);
-#endif // defined(AMD_MFMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA4))
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_128B_AVAILABLE)
     }
 
 #if defined(TURING_MMA_AVAILABLE)
@@ -726,7 +726,7 @@ namespace ggml_cuda_mma {
 
         return ret;
     }
-#elif defined(AMD_WMMA_AVAILABLE) && defined(RDNA3)
+#elif defined(AMD_WMMA_256B_AVAILABLE)
     static __device__ __forceinline__ tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> get_half2(
             const tile<16, 16, float, DATA_LAYOUT_I_MAJOR> & tile_float) {
         tile<16, 8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> ret;
@@ -1171,7 +1171,7 @@ namespace ggml_cuda_mma {
             : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[3]));
 #endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
 #elif defined(AMD_WMMA_AVAILABLE)
-#if defined(RDNA4)
+#if defined(AMD_WMMA_128B_AVAILABLE)
         using halfx8_t = __attribute__((ext_vector_type(8))) _Float16;
         halfx8_t& acc_frag = reinterpret_cast<halfx8_t&>(D.x[0]);
         const halfx8_t& a_frag = reinterpret_cast<const halfx8_t&>(A.x[0]);
@@ -1180,7 +1180,7 @@ namespace ggml_cuda_mma {
 #else
         GGML_UNUSED_VARS(D, A, B);
         NO_DEVICE_CODE;
-#endif // defined(RDNA4)
+#endif // defined(AMD_WMMA_128B_AVAILABLE)
 #elif defined(AMD_MFMA_AVAILABLE)
         // MFMA: FP16 input, FP32 accumulate, convert back to half2.
         using halfx4_t = __attribute__((ext_vector_type(4))) _Float16;
@@ -1218,7 +1218,7 @@ namespace ggml_cuda_mma {
     static __device__ __forceinline__ void mma(
             tile<16, 16, half2, DATA_LAYOUT_I_MAJOR_SCRAMBLED> & D, const tile<32, 8, half2, DATA_LAYOUT_I_MAJOR> & A,
             const tile<16, 8, half2, DATA_LAYOUT_I_MAJOR> & B) {
-#if defined(AMD_MFMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA4))
+#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_128B_AVAILABLE)
         tile<16, 8, half2>       * D16 = (tile<16, 8, half2>       *) &D;
         const tile<16, 8, half2> * A16 = (const tile<16, 8, half2> *) &A;
         mma(D16[0], A16[0], B);
@@ -1226,7 +1226,7 @@ namespace ggml_cuda_mma {
 #else
         GGML_UNUSED_VARS(D, A, B);
         NO_DEVICE_CODE;
-#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) && defined(RDNA4)
+#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_128B_AVAILABLE)
     }
 
     template <data_layout dl_ab, data_layout dl_d>
@@ -1373,14 +1373,14 @@ namespace ggml_cuda_mma {
             : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[3]));
 #endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE
 #elif defined(AMD_WMMA_AVAILABLE)
-#if defined(RDNA4)
+#if defined(AMD_WMMA_128B_AVAILABLE)
         using halfx8_t = __attribute__((ext_vector_type(8))) _Float16;
         using floatx8_t = __attribute__((ext_vector_type(8))) float;
         floatx8_t& acc_frag = reinterpret_cast<floatx8_t&>(D.x[0]);
         const halfx8_t& a_frag = reinterpret_cast<const halfx8_t&>(A.x[0]);
         const halfx8_t& b_frag = reinterpret_cast<const halfx8_t&>(B.x[0]);
         acc_frag = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(a_frag, b_frag, acc_frag);
-#elif defined(RDNA3)
+#elif defined(AMD_WMMA_256B_AVAILABLE)
         using halfx16_t = __attribute__((ext_vector_type(16))) _Float16;
         using floatx8_t = __attribute__((ext_vector_type(8))) float;
         floatx8_t& acc_frag = reinterpret_cast<floatx8_t&>(D.x[0]);
@@ -1390,7 +1390,7 @@ namespace ggml_cuda_mma {
 #else
         GGML_UNUSED_VARS(D, A, B);
         NO_DEVICE_CODE;
-#endif // RDNA4
+#endif // AMD_WMMA_128B_AVAILABLE
 #elif defined(AMD_MFMA_AVAILABLE)
         using halfx4_t = __attribute__((ext_vector_type(4))) _Float16;
         using floatx4_t = __attribute__((ext_vector_type(4))) float;
@@ -1408,14 +1408,14 @@ namespace ggml_cuda_mma {
     static __device__ __forceinline__ void mma(
             tile<16, 16, float, dl_d> & D, const tile<16, 8, nv_bfloat162, dl_ab> & A, const tile<16, 8, nv_bfloat162, dl_ab> & B) {
 #if defined(AMD_WMMA_AVAILABLE)
-#if defined(RDNA4)
+#if defined(AMD_WMMA_128B_AVAILABLE)
         using bf16x8_t = __attribute__((ext_vector_type(8))) short;
         using floatx8_t = __attribute__((ext_vector_type(8))) float;
         floatx8_t& acc_frag = reinterpret_cast<floatx8_t&>(D.x[0]);
         const bf16x8_t& a_frag = reinterpret_cast<const bf16x8_t&>(A.x[0]);
         const bf16x8_t& b_frag = reinterpret_cast<const bf16x8_t&>(B.x[0]);
         acc_frag = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32_gfx12(a_frag, b_frag, acc_frag);
-#elif defined(RDNA3)
+#elif defined(AMD_WMMA_256B_AVAILABLE)
         using bf16x16_t = __attribute__((ext_vector_type(16))) short;
         using floatx8_t = __attribute__((ext_vector_type(8))) float;
         floatx8_t& acc_frag = reinterpret_cast<floatx8_t&>(D.x[0]);
@@ -1425,7 +1425,7 @@ namespace ggml_cuda_mma {
 #else
         GGML_UNUSED_VARS(D, A, B);
         NO_DEVICE_CODE;
-#endif // defined(RDNA4)
+#endif // defined(AMD_WMMA_128B_AVAILABLE)
 #elif defined(AMD_MFMA_AVAILABLE)
         using floatx4_t = __attribute__((ext_vector_type(4))) float;
         floatx4_t& acc_frag = reinterpret_cast<floatx4_t&>(D.x[0]);
@@ -1467,19 +1467,19 @@ namespace ggml_cuda_mma {
 #elif defined(AMD_WMMA_AVAILABLE)
         using int32x8_t = __attribute__((__vector_size__(8 * sizeof(int)))) int;
         int32x8_t * acc = (int32x8_t *) D.x;
-#if defined(RDNA4)
+#if defined(AMD_WMMA_128B_AVAILABLE)
         using int32x2_t = __attribute__((__vector_size__(2 * sizeof(int)))) int;
         int32x2_t * a_vec = (int32x2_t *) A.x;
         int32x2_t * b_vec = (int32x2_t *) B.x;
         acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32_gfx12(true, a_vec[0], true, b_vec[0], acc[0], true);
         acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32_gfx12(true, a_vec[1], true, b_vec[1], acc[0], true);
-#elif defined(RDNA3)
+#elif defined(AMD_WMMA_256B_AVAILABLE)
         using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int;
         int32x4_t * a_vec = (int32x4_t *) A.x;
         int32x4_t * b_vec = (int32x4_t *) B.x;
         acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, a_vec[0], true, b_vec[0], acc[0], true);
         acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, a_vec[1], true, b_vec[1], acc[0], true);
-#endif // RDNA4
+#endif // AMD_WMMA_128B_AVAILABLE
 #else
         GGML_UNUSED_VARS(D, A, B);
         NO_DEVICE_CODE;
@@ -1556,7 +1556,7 @@ namespace ggml_cuda_mma {
     static __device__ __forceinline__ void mma(
             tile<16, 16, half2, DATA_LAYOUT_I_MAJOR> & D, const tile<32,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> & A,
             const tile<16,  8, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> & B) {
-#if defined(AMD_WMMA_AVAILABLE) && defined(RDNA3)
+#if defined(AMD_WMMA_256B_AVAILABLE)
         using halfx16_t = __attribute__((ext_vector_type(16))) _Float16;
         halfx16_t       * xD = (halfx16_t       *) D.x;
         const halfx16_t * xA = (const halfx16_t *) A.x;
@@ -1585,17 +1585,17 @@ namespace ggml_cuda_mma {
 #elif defined(AMD_WMMA_AVAILABLE)
         using int32x8_t = __attribute__((__vector_size__(8 * sizeof(int)))) int;
         int32x8_t * acc = (int32x8_t *) D.x;
-#if defined(RDNA4)
+#if defined(AMD_WMMA_128B_AVAILABLE)
         using int32x2_t = __attribute__((__vector_size__(2 * sizeof(int)))) int;
         int32x2_t * a_vec = (int32x2_t *) A.x;
         int32x2_t * b_vec = (int32x2_t *) B.x;
         acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32_gfx12(true, a_vec[0], true, b_vec[0], acc[0], false);
-#elif defined(RDNA3)
+#elif defined(AMD_WMMA_256B_AVAILABLE)
         using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int;
         int32x4_t * a_vec = (int32x4_t *) A.x;
         int32x4_t * b_vec = (int32x4_t *) B.x;
         acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, a_vec[0], true, b_vec[0], acc[0], false);
-#endif // RDNA4
+#endif // AMD_WMMA_128B_AVAILABLE
 #else
         GGML_UNUSED(D);
         GGML_UNUSED(A);

ggml/src/ggml-cuda/mmq.cu

@@ -242,7 +242,7 @@ void ggml_cuda_mul_mat_q(
 
             // Each expert only sees ne12*n_expert_used/ne02 tokens on average.
             // On RDNA3 and RDNA4 it is faster to pick the tile size against this value instead of ne12.
-            if (GGML_CUDA_CC_IS_RDNA3(cc) || GGML_CUDA_CC_IS_RDNA4(cc)) {
+            if (GGML_CUDA_CC_IS_RDNA3(cc) || GGML_CUDA_CC_IS_RDNA4_M(cc) || GGML_CUDA_CC_IS_RDNA4(cc)) {
                 ncols_opt = (ne12*n_expert_used + ne02 - 1) / ne02;
             }
         }
@@ -470,7 +470,7 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t
     }
 
     if (amd_wmma_available(cc)) {
-        if (GGML_CUDA_CC_IS_RDNA3(cc)) {
+        if (GGML_CUDA_CC_IS_RDNA3(cc) || GGML_CUDA_CC_IS_RDNA4_M(cc)) {
             // High expert counts are almost always better on MMQ due to
             //     the synchronization overhead in the cuBLAS/hipBLAS path:
             // https://github.com/ggml-org/llama.cpp/pull/18202
@@ -487,7 +487,7 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t
                     return ne11 <= (GGML_CUDA_CC_IS_RDNA3_0(cc) ? 128 : 256);
                 case GGML_TYPE_IQ2_XS:
                 case GGML_TYPE_IQ2_S:
-                    return GGML_CUDA_CC_IS_RDNA3_5(cc) || ne11 <= 128;
+                    return GGML_CUDA_CC_IS_RDNA3_5(cc) || GGML_CUDA_CC_IS_RDNA4_M(cc) || ne11 <= 128;
                 default:
                     return true;
             }

ggml/src/ggml-cuda/mmq.cuh

@@ -240,6 +240,9 @@ static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type ty
         if (GGML_CUDA_CC_IS_RDNA4(cc)) {
             return ggml_cuda_mmq_get_config_rdna4(type, J, fallback);
         }
+        if (GGML_CUDA_CC_IS_RDNA4_M(cc)) {
+            return ggml_cuda_mmq_get_config_rdna3_5(type, J, fallback);
+        }
         if (GGML_CUDA_CC_IS_RDNA3_5(cc)) {
             return ggml_cuda_mmq_get_config_rdna3_5(type, J, fallback);
         }
@@ -272,6 +275,8 @@ static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_t
     return ggml_cuda_mmq_get_config_cdna(type, J, fallback);
 #elif defined(RDNA4)
     return ggml_cuda_mmq_get_config_rdna4(type, J, fallback);
+#elif defined(RDNA4_M)
+    return ggml_cuda_mmq_get_config_rdna3_5(type, J, fallback);
 #elif defined(RDNA3_5)
     return ggml_cuda_mmq_get_config_rdna3_5(type, J, fallback);
 #elif defined(RDNA3)

ggml/src/ggml-cuda/mmvf.cu

@@ -844,7 +844,7 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, int warp_size, const
                     return ne11 <= MMVF_MAX_BATCH_SIZE;
                 }
                 if (fp16_mma_hardware_available(cc)) {
-                    if (GGML_CUDA_CC_IS_RDNA3(cc)) {
+                    if (GGML_CUDA_CC_IS_RDNA3(cc) || GGML_CUDA_CC_IS_RDNA4_M(cc)) {
                         return ne11 <= 3;
                     }
                     if (GGML_CUDA_CC_IS_RDNA4(cc)) {

ggml/src/ggml-cuda/mmvq.cu

@@ -108,7 +108,7 @@ static constexpr __device__ mmvq_parameter_table_id get_device_table_id() {
     return MMVQ_PARAMETERS_RDNA4;
 #elif defined(RDNA3_0)
     return MMVQ_PARAMETERS_RDNA3_0;
-#elif defined(RDNA2) || defined(RDNA3_5)
+#elif defined(RDNA2) || defined(RDNA3_5) || defined(RDNA4_M)
     return MMVQ_PARAMETERS_RDNA2;
 #elif defined(GCN) || defined(CDNA)
     return MMVQ_PARAMETERS_GCN;
@@ -128,7 +128,7 @@ static __host__ mmvq_parameter_table_id get_device_table_id(int cc) {
     if (GGML_CUDA_CC_IS_RDNA3_0(cc)) {
         return MMVQ_PARAMETERS_RDNA3_0;
     }
-    if (GGML_CUDA_CC_IS_RDNA2(cc) || GGML_CUDA_CC_IS_RDNA3_5(cc)) {
+    if (GGML_CUDA_CC_IS_RDNA2(cc) || GGML_CUDA_CC_IS_RDNA3_5(cc) || GGML_CUDA_CC_IS_RDNA4_M(cc)) {
         return MMVQ_PARAMETERS_RDNA2;
     }
     if (GGML_CUDA_CC_IS_GCN(cc) || GGML_CUDA_CC_IS_CDNA(cc)) {
@@ -299,7 +299,7 @@ int get_mmvq_mmid_max_batch(ggml_type type, int cc) {
         if (GGML_CUDA_CC_IS_RDNA4(cc)) {
             return get_mmvq_mmid_max_batch_rdna4(type);
         }
-        if (GGML_CUDA_CC_IS_RDNA3(cc)) {
+        if (GGML_CUDA_CC_IS_RDNA3(cc) || GGML_CUDA_CC_IS_RDNA4_M(cc)) {
             return get_mmvq_mmid_max_batch_rdna3(type);
         }
         if (GGML_CUDA_CC_IS_RDNA1(cc) || GGML_CUDA_CC_IS_RDNA2(cc)) {
@@ -432,7 +432,7 @@ template <ggml_type type>
 static constexpr __device__ int get_mmvq_mmid_max_batch_for_device() {
 #if defined(RDNA4)
     return get_mmvq_mmid_max_batch_rdna4(type);
-#elif defined(RDNA3)
+#elif defined(RDNA3) || defined(RDNA4_M)
     return get_mmvq_mmid_max_batch_rdna3(type);
 #elif defined(RDNA2) || defined(RDNA1)
     return get_mmvq_mmid_max_batch_rdna1_rdna2(type);

ggml/src/ggml-cuda/quantize.cu

@@ -83,21 +83,21 @@ static __global__ void quantize_q8_1(
     ggml_cuda_pdl_sync();
     const float xi = i0 < ne00 ? x[i03*s03 + i02*s02 + i01*s01 + i00] : 0.0f;
     float amax = fabsf(xi);
-    float sum = xi;
 
     amax = warp_reduce_max<QK8_1>(amax);
-    sum  = warp_reduce_sum<QK8_1>(sum);
 
     const float  d = amax / 127.0f;
     const int8_t q = amax == 0.0f ? 0 : roundf(xi / d);
 
+    const int sum = warp_reduce_sum<QK8_1>((int) q);
+
     y[ib].qs[iqs] = q;
 
     if (iqs > 0) {
         return;
     }
 
-    y[ib].ds = make_half2(d, sum);
+    y[ib].ds = make_half2(d, d * sum);
 }
 
 __device__ __forceinline__ uint8_t compute_e8m0_scale(float amax) {
@@ -500,17 +500,6 @@ static __global__ void quantize_mmq_q8_1(
         amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, offset, WARP_SIZE));
     }
 
-    float sum;
-    if (ds_layout != MMQ_Q8_1_DS_LAYOUT_D4) {
-        sum = xi.x + xi.y + xi.z + xi.w;
-
-        // Calculate sums across vals_per_sum/4 threads.
-#pragma unroll
-        for (int offset = vals_per_sum/8; offset > 0; offset >>= 1) {
-            sum += __shfl_xor_sync(0xFFFFFFFF, sum, offset, WARP_SIZE);
-        }
-    }
-
     const float d_inv = 127.0f / amax;
     char4 q;
     q.x = roundf(xi.x*d_inv);
@@ -519,6 +508,19 @@ static __global__ void quantize_mmq_q8_1(
     q.w = roundf(xi.w*d_inv);
     const float d = 1.0f / d_inv;
 
+    float sum = 0.0f;
+    if (ds_layout != MMQ_Q8_1_DS_LAYOUT_D4) {
+        int sum_q = int(q.x) + int(q.y) + int(q.z) + int(q.w);
+
+        // Calculate sums across vals_per_sum/4 threads.
+#pragma unroll
+        for (int offset = vals_per_sum/8; offset > 0; offset >>= 1) {
+            sum_q += __shfl_xor_sync(0xFFFFFFFF, sum_q, offset, WARP_SIZE);
+        }
+
+        sum = d * float(sum_q);
+    }
+
     // write the block once (normal) or to each of the token's compact rows (scatter)
     const int nwrite = scatter ? n_expert_used : 1;
 #pragma unroll

ggml/src/ggml-cuda/vendors/hip.h

@@ -216,9 +216,13 @@
 #define RDNA4
 #endif // defined(__GFX12__)
 
-#if defined(__GFX11__)
+#if defined(__gfx1170__) || defined(__gfx1171__)
+#define RDNA4_M
+#endif // defined(__gfx1170__) || defined(__gfx1171__)
+
+#if defined(__GFX11__) && !defined(RDNA4_M)
 #define RDNA3
-#endif // defined(__GFX11__)
+#endif // defined(__GFX11__) && !defined(RDNA4_M)
 
 #if defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__)
 #define RDNA3_5
@@ -237,9 +241,9 @@
 #define RDNA1
 #endif // defined(__gfx1010__) || defined(__gfx1012__)
 
-#if defined(RDNA4) || defined(RDNA3) || defined(RDNA2) || defined(RDNA1)
+#if defined(RDNA4) || defined(RDNA4_M) || defined(RDNA3) || defined(RDNA2) || defined(RDNA1)
 #define RDNA // For the entire family
-#endif // defined(RDNA4) || defined(RDNA3) || defined(RDNA2) || defined(RDNA1)
+#endif // defined(RDNA4) || defined(RDNA4_M) || defined(RDNA3) || defined(RDNA2) || defined(RDNA1)
 
 #ifndef __has_builtin
     #define __has_builtin(x) 0
@@ -254,6 +258,10 @@ typedef __hip_fp8_e4m3 __nv_fp8_e4m3;
 #define FP8_AVAILABLE
 #endif // HIP_VERSION >= 60300000
 
+#if HIP_VERSION < 70000000
+#define __ballot_sync(mask, var) __ballot(var)
+#endif //HIP_VERSION < 70000000
+
 typedef int8_t int8x4_t __attribute__((ext_vector_type(4)));
 typedef uint8_t uint8x4_t __attribute__((ext_vector_type(4)));
 static __device__ __forceinline__ int __vsubss4(const int a, const int b) {

ggml/src/ggml-et/ggml-et.cpp

@@ -373,6 +373,9 @@ static const struct ggml_backend_buffer_i ggml_backend_et_buffer_i = {
     /* .cpy_tensor      = */ ggml_backend_et_buffer_cpy_tensor,
     /* .clear           = */ ggml_backend_et_buffer_clear,
     /* .reset           = */ NULL,
+    /* .set_tensor_n    = */ NULL,
+    /* .get_tensor_n    = */ NULL,
+    /* .cpy_tensor_n    = */ NULL,
 };
 
 static const char * ggml_backend_et_buffer_type_get_name(ggml_backend_buffer_type_t buft) {
@@ -1270,7 +1273,7 @@ static bool ggml_backend_et_device_supports_op(ggml_backend_dev_t dev, const ggm
         case GGML_OP_FLASH_ATTN_EXT:
             if (op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 && op->src[1] &&
                 (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) && op->src[2] &&
-                (op->src[2]->type == GGML_TYPE_F32 || op->src[2]->type == GGML_TYPE_F16) && op->src[4] == nullptr &&
+                (op->src[2]->type == GGML_TYPE_F32 || op->src[2]->type == GGML_TYPE_F16) && op->src[4] == nullptr && op->src[5] == nullptr &&
                 ggml_is_contiguous_rows(op) && ggml_is_contiguous_rows(op->src[0])) {
                 float max_bias      = 0.0f;
                 float logit_softcap = 0.0f;

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

@@ -843,6 +843,9 @@ static ggml_backend_buffer_i ggml_backend_hexagon_buffer_interface = {
     /* .cpy_tensor      = */ ggml_backend_hexagon_buffer_cpy_tensor,
     /* .clear           = */ ggml_backend_hexagon_buffer_clear,
     /* .reset           = */ NULL,
+    /* .set_tensor_n    = */ NULL,
+    /* .get_tensor_n    = */ NULL,
+    /* .cpy_tensor_n    = */ NULL,
 };
 
 // ** backend buffer type
@@ -877,6 +880,9 @@ static ggml_backend_buffer_i ggml_backend_hexagon_host_buffer_interface = {
     /* .cpy_tensor      = */ ggml_backend_hexagon_buffer_cpy_tensor,
     /* .clear           = */ ggml_backend_hexagon_buffer_clear,
     /* .reset           = */ NULL,
+    /* .set_tensor_n    = */ NULL,
+    /* .get_tensor_n    = */ NULL,
+    /* .cpy_tensor_n    = */ NULL,
 };
 
 // ** backend buffer type

ggml/src/ggml-hexagon/htp-kparams.cpp

@@ -60,6 +60,7 @@ htp_op_code htp_op_remap(const struct ggml_tensor * t) {
         case GGML_OP_CONV_2D:         return HTP_OP_CONV_2D;
         case GGML_OP_CONV_2D_DW:      return HTP_OP_CONV_2D_DW;
         case GGML_OP_ROLL:            return HTP_OP_ROLL;
+        case GGML_OP_UPSCALE:         return HTP_OP_UPSCALE;
 
         case GGML_OP_UNARY:
             switch (ggml_get_unary_op(t)) {
@@ -347,6 +348,11 @@ static bool htp_kparams_solve_flash_attn(
         return false;
     }
 
+    // K/V row lists (ggml_flash_attn_ext_rows) are not supported
+    if (op->src[5] != nullptr) {
+        return false;
+    }
+
     memset(kparams, 0, sizeof(*kparams));
 
     const struct ggml_tensor * q     = op->src[0];
@@ -2185,7 +2191,8 @@ static bool htp_kparams_solve_ssm_conv(
         target_max_tile          = (std::max)(target_max_tile, 32u);
         target_max_tile          = (std::min)(target_max_tile, 128u);
 
-        uint32_t d_inner_tile = (uint32_t) (avail_for_tiles / (2 * (ncs + n_t) * sizeof(float)));
+        const size_t raw_stride_bytes = ncs * sizeof(float) + sizeof(float);
+        uint32_t d_inner_tile = (uint32_t) (avail_for_tiles / (2 * (raw_stride_bytes + n_t * sizeof(float))));
         d_inner_tile          = (d_inner_tile / 32) * 32;
         if (d_inner_tile == 0) {
             d_inner_tile = 32;
@@ -2372,6 +2379,183 @@ static bool htp_kparams_solve_pool_2d(
     return true;
 }
 
+bool htp_kparams_solve_upscale(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * dst,
+    struct htp_upscale_kernel_params * kparams
+) {
+    memset(kparams, 0, sizeof(*kparams));
+
+    if (src0->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+        return false;
+    }
+
+    const uint32_t src_x = (uint32_t) src0->ne[0];
+    const uint32_t src_y = (uint32_t) src0->ne[1];
+    const uint32_t src_z = (uint32_t) src0->ne[2];
+    const uint32_t src_w = (uint32_t) src0->ne[3];
+
+    const uint32_t dst_x = (uint32_t) dst->ne[0];
+    const uint32_t dst_y = (uint32_t) dst->ne[1];
+    const uint32_t dst_z = (uint32_t) dst->ne[2];
+    const uint32_t dst_w = (uint32_t) dst->ne[3];
+
+    if (src_x == 0 || src_y == 0 || src_z == 0 || src_w == 0 ||
+        dst_x == 0 || dst_y == 0 || dst_z == 0 || dst_w == 0) {
+        return false;
+    }
+
+    const bool is_src_contig_row = (src0->nb[0] == sizeof(float));
+    const bool is_src_strided_2d = (src0->nb[0] > sizeof(float)) &&
+                                   (src0->nb[0] % sizeof(float) == 0) &&
+                                   (src0->nb[1] == (size_t) src_x * src0->nb[0]) &&
+                                   ((size_t) src_x * src_y <= 0xFFFFu);
+
+    if ((!is_src_contig_row && !is_src_strided_2d) || dst->nb[0] != sizeof(float)) {
+        return false;
+    }
+
+    const int32_t mode_flags = dst->op_params[0];
+    const uint32_t mode_raw  = (uint32_t) (mode_flags & 0xFF);
+    if (mode_raw >= GGML_SCALE_MODE_COUNT) {
+        return false;
+    }
+
+    const bool align_corners = (mode_flags & GGML_SCALE_FLAG_ALIGN_CORNERS) != 0;
+    const bool antialias     = (mode_flags & GGML_SCALE_FLAG_ANTIALIAS) != 0;
+
+    float sf0 = (float) dst_x / (float) src_x;
+    float sf1 = (float) dst_y / (float) src_y;
+    float sf2 = (float) dst_z / (float) src_z;
+    float sf3 = (float) dst_w / (float) src_w;
+    float pixel_offset = 0.5f;
+
+    if (align_corners) {
+        pixel_offset = 0.0f;
+        sf0 = (dst_x > 1 && src_x > 1) ? (float) (dst_x - 1) / (float) (src_x - 1) : sf0;
+        sf1 = (dst_y > 1 && src_y > 1) ? (float) (dst_y - 1) / (float) (src_y - 1) : sf1;
+    }
+
+    const float inv_sf0 = 1.0f / sf0;
+    const float inv_sf1 = 1.0f / sf1;
+    const float inv_sf2 = 1.0f / sf2;
+    const float inv_sf3 = 1.0f / sf3;
+
+    const float filter_radius0 = antialias ? (std::max)(1.0f, inv_sf0) : 0.0f;
+    const float filter_radius1 = antialias ? (std::max)(1.0f, inv_sf1) : 0.0f;
+    const float invscale0      = antialias ? (1.0f / filter_radius0) : 0.0f;
+    const float invscale1      = antialias ? (1.0f / filter_radius1) : 0.0f;
+
+    const uint32_t total_slices = dst_z * dst_w;
+    uint32_t       n_threads    = htp_session_n_threads(sess);
+
+    struct htp_upscale_vtcm_layout layout;
+    uint32_t dst_nrows_per_chunk = 0;
+    while (n_threads > 0) {
+        if (htp_upscale_solve_layout(&layout, src_x, src_y, dst_x, dst_y, inv_sf1, filter_radius1, n_threads,
+                                     htp_session_vtcm_size(sess), &dst_nrows_per_chunk)) {
+            break;
+        }
+        n_threads--;
+    }
+    if (n_threads == 0) {
+        return false;
+    }
+
+    size_t lut_x_size = 0;
+    size_t elem_y_sz  = sizeof(uint32_t);
+    if (mode_raw == HTP_UPSCALE_MODE_BICUBIC) {
+        lut_x_size = hex_round_up(dst_x * sizeof(struct upscale_bicubic_x), 128);
+        elem_y_sz  = sizeof(struct upscale_bicubic_y);
+    } else if (antialias) {
+        lut_x_size = hex_round_up(dst_x * sizeof(struct upscale_antialias_x), 128);
+        elem_y_sz  = sizeof(struct upscale_antialias_y);
+    } else if (mode_raw == HTP_UPSCALE_MODE_BILINEAR) {
+        const uint32_t n_chunks = (dst_x + 31) / 32;
+        lut_x_size = hex_round_up(n_chunks * sizeof(struct upscale_bilinear_x_chunk), 128);
+        elem_y_sz  = sizeof(struct upscale_bilinear_y);
+    } else {
+        lut_x_size = hex_round_up(dst_x * sizeof(uint32_t), 128);
+    }
+    const size_t total_lut_size = lut_x_size + hex_round_up(dst_y * elem_y_sz, 128);
+    if (total_lut_size > HTP_DDR_SPAD_SIZE) {
+        return false;
+    }
+
+    const uint32_t chunks_per_slice = (dst_y + dst_nrows_per_chunk - 1) / dst_nrows_per_chunk;
+    const uint32_t total_chunks     = total_slices * chunks_per_slice;
+
+    n_threads = (std::min)(n_threads, total_chunks > 0 ? total_chunks : 1);
+    htp_upscale_vtcm_layout_build(&layout, src_x, src_y, dst_x, dst_y, dst_nrows_per_chunk, inv_sf1, filter_radius1, n_threads);
+
+    const uint32_t slice_dst_bytes = dst_x * dst_y * sizeof(float);
+    const uint32_t chunk_dst_bytes = dst_x * dst_nrows_per_chunk * sizeof(float);
+    const bool     dst_contig      = (dst->nb[1] == (size_t) dst_x * sizeof(float)) &&
+                                     (dst_z <= 1 || dst->nb[2] == (size_t) dst_y * dst->nb[1]) &&
+                                     (dst_w <= 1 || dst->nb[3] == (size_t) dst_z * dst->nb[2]);
+
+    uint32_t chunks_per_group = 0;
+    if (dst_contig) {
+        const uint32_t slices_per_group = (slice_dst_bytes > 0)
+            ? (HEX_L2_LINE_SIZE / hex_gcd_u32(slice_dst_bytes, HEX_L2_LINE_SIZE))
+            : 1;
+        if (slices_per_group == 1 && (chunk_dst_bytes % HEX_L2_LINE_SIZE == 0)) {
+            chunks_per_group = 1;
+        } else {
+            chunks_per_group = slices_per_group * chunks_per_slice;
+        }
+    }
+
+    uint8_t kernel_type = HTP_UPSCALE_KERNEL_BILINEAR;
+    if (mode_raw == HTP_UPSCALE_MODE_NEAREST) {
+        if (dst_x == 2 * src_x && dst_y == 2 * src_y && !align_corners) {
+            kernel_type = HTP_UPSCALE_KERNEL_NEAREST_2X;
+        } else if (dst_x == 4 * src_x && dst_y == 4 * src_y && !align_corners) {
+            kernel_type = HTP_UPSCALE_KERNEL_NEAREST_4X;
+        } else {
+            kernel_type = HTP_UPSCALE_KERNEL_NEAREST;
+        }
+    } else if (mode_raw == HTP_UPSCALE_MODE_BICUBIC) {
+        kernel_type = HTP_UPSCALE_KERNEL_BICUBIC;
+    } else if (antialias) {
+        kernel_type = HTP_UPSCALE_KERNEL_BILINEAR_AA;
+    } else if (dst_x == 2 * src_x && dst_y == 2 * src_y && !align_corners && (src_x % 32 == 0)) {
+        kernel_type = HTP_UPSCALE_KERNEL_BILINEAR_2X;
+    } else {
+        kernel_type = HTP_UPSCALE_KERNEL_BILINEAR;
+    }
+
+    kparams->kernel_type          = kernel_type;
+    kparams->mode                 = (uint8_t) mode_raw;
+    kparams->align_corners        = align_corners ? 1 : 0;
+    kparams->antialias            = antialias ? 1 : 0;
+
+    kparams->inv_sf0              = inv_sf0;
+    kparams->inv_sf1              = inv_sf1;
+    kparams->inv_sf2              = inv_sf2;
+    kparams->inv_sf3              = inv_sf3;
+    kparams->pixel_offset         = pixel_offset;
+
+    kparams->filter_radius0       = filter_radius0;
+    kparams->filter_radius1       = filter_radius1;
+    kparams->invscale0            = invscale0;
+    kparams->invscale1            = invscale1;
+
+    kparams->n_threads            = n_threads;
+    kparams->chunks_per_group     = chunks_per_group;
+
+    kparams->dst_nrows_per_chunk  = dst_nrows_per_chunk;
+    kparams->chunks_per_slice     = chunks_per_slice;
+    kparams->total_chunks         = total_chunks;
+    kparams->vtcm_size            = (uint32_t) layout.total_bytes;
+
+    kparams->div_chunks_per_slice = init_fastdiv_values(chunks_per_slice > 0 ? chunks_per_slice : 1);
+    kparams->div_dst_z            = init_fastdiv_values(dst_z > 0 ? dst_z : 1);
+
+    return true;
+}
+
 static bool htp_kparams_solve_concat(
     const struct ggml_hexagon_session * sess,
     const struct ggml_tensor * op,
@@ -2476,6 +2660,129 @@ static bool htp_kparams_solve_concat(
     return false;
 }
 
+// Permuted copy with strided dim 0 (e.g. channel-first <-> channel-last)
+static bool htp_kparams_solve_cpy_transpose(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * dst,
+    uint32_t type_size,
+    struct htp_copy_kernel_params * kparams) {
+
+    // View contiguous dst with src shape to support CPY with reshape
+    const bool dst_cont = ggml_is_contiguous(dst);
+    if (!dst_cont && !ggml_are_same_shape(src0, dst)) {
+        return false;
+    }
+
+    // Drop size-1 dims and merge dims adjacent in both tensors
+    size_t ne[GGML_MAX_DIMS];
+    size_t nbs[GGML_MAX_DIMS];
+    size_t nbd[GGML_MAX_DIMS];
+    size_t dst_stride = type_size;
+    int    n          = 0;
+    for (int d = 0; d < GGML_MAX_DIMS; d++) {
+        const size_t sd = dst_cont ? dst_stride : dst->nb[d];
+        dst_stride *= (size_t) src0->ne[d];
+        if (src0->ne[d] == 1) {
+            continue;
+        }
+        if (n > 0 && src0->nb[d] == nbs[n - 1] * ne[n - 1] && sd == nbd[n - 1] * ne[n - 1]) {
+            ne[n - 1] *= (size_t) src0->ne[d];
+            continue;
+        }
+        ne[n]  = (size_t) src0->ne[d];
+        nbs[n] = src0->nb[d];
+        nbd[n] = sd;
+        n++;
+    }
+
+    // Fold leading dense dim into element block
+    size_t es = type_size;
+    for (int i = 0; i < n; i++) {
+        if (nbs[i] == type_size && nbd[i] == type_size) {
+            es *= ne[i];
+            for (int j = i; j + 1 < n; j++) {
+                ne[j]  = ne[j + 1];
+                nbs[j] = nbs[j + 1];
+                nbd[j] = nbd[j + 1];
+            }
+            n--;
+            break;
+        }
+    }
+    if (es > 128 || (es & (es - 1)) != 0) {
+        return false;
+    }
+
+    // Transpose axes: dim a dense in dst, dim k dense in src
+    int a = -1;
+    int k = -1;
+    for (int i = 0; i < n; i++) {
+        if (a < 0 && nbd[i] == es) {
+            a = i;
+        }
+        if (k < 0 && nbs[i] == es) {
+            k = i;
+        }
+    }
+    if (a < 0 || k < 0 || a == k) {
+        return false;
+    }
+
+    struct htp_copy_transpose_params tp = {};
+    tp.elem_size = (uint8_t) es;
+    tp.ne[0]     = (uint32_t) ne[a];
+    tp.ne[1]     = (uint32_t) ne[k];
+    tp.ne[2]     = 1;
+    tp.src_nb[0] = (uint32_t) nbs[a];
+    tp.dst_nb[0] = (uint32_t) nbd[k];
+    tp.packed    = nbd[k] == ne[a] * es && ne[a] * es < 128 && 128 % (ne[a] * es) == 0;
+
+    size_t ne3 = 1;
+    int    m   = 0;
+    for (int i = 0; i < n; i++) {
+        if (i == a || i == k) {
+            continue;
+        }
+        if (m == 0) {
+            tp.ne[2] = (uint32_t) ne[i];
+        } else {
+            ne3 = ne[i];
+        }
+        tp.src_nb[1 + m] = (uint32_t) nbs[i];
+        tp.dst_nb[1 + m] = (uint32_t) nbd[i];
+        m++;
+    }
+
+    const uint32_t n_threads = htp_session_n_threads(sess);
+    const size_t   vtcm_size = htp_session_vtcm_size(sess);
+    uint64_t       vtcm      = 0;
+    for (uint32_t blk_bytes = HTP_COPY_TRANSPOSE_BLK_BYTES; ; blk_bytes /= 2) {
+        if (!htp_copy_transpose_layout_build(&tp, blk_bytes)) {
+            return false;
+        }
+        vtcm = (uint64_t) n_threads * 2 * (tp.in_buf_size + tp.out_buf_size);
+        if (vtcm <= vtcm_size) {
+            break;
+        }
+    }
+
+    const uint64_t total_tiles = (uint64_t) tp.n0b * tp.nkb * tp.ne[2] * ne3;
+    if (total_tiles > UINT32_MAX) {
+        return false;
+    }
+    tp.total_tiles = (uint32_t) total_tiles;
+    tp.div_nkb     = init_fastdiv_values(tp.nkb);
+    tp.div_n0b     = init_fastdiv_values(tp.n0b);
+    tp.div_ne2     = init_fastdiv_values(tp.ne[2]);
+
+    kparams->kernel_type = HTP_COPY_KERNEL_TRANSPOSE;
+    kparams->n_threads   = (uint8_t) n_threads;
+    kparams->vtcm_size   = (uint32_t) vtcm;
+    kparams->u.transpose = tp;
+    return true;
+}
+
 static bool htp_kparams_solve_cpy(
     const struct ggml_hexagon_session * sess,
     const struct ggml_tensor * op,
@@ -2573,6 +2880,10 @@ static bool htp_kparams_solve_cpy(
             }
         }
 
+        if (htp_kparams_solve_cpy_transpose(sess, src0, dst, src_type_size, kparams)) {
+            return true;
+        }
+
         if (sameshape) {
             kparams->kernel_type = HTP_COPY_KERNEL_SAMESHAPE_SAMETYPE;
             kparams->total_rows  = (uint32_t) (src0->ne[1] * src0->ne[2] * src0->ne[3]);
@@ -3916,6 +4227,31 @@ void htp_opnode_collapse(const struct ggml_hexagon_session * sess, htp_opnode &
         }
         return;
     }
+
+    if (node.opcode == HTP_OP_UPSCALE) {
+        if (node.inputs.empty() || !node.inputs[0] || node.outputs.empty() || !node.outputs[0]) {
+            return;
+        }
+        const struct ggml_tensor * src0 = node.inputs[0];
+        const struct ggml_tensor * dst  = node.outputs[0];
+
+        if ((src0->ne[3] > 1 || dst->ne[3] > 1) &&
+            src0->ne[2] == dst->ne[2] && src0->ne[3] == dst->ne[3] &&
+            src0->nb[3] == (size_t) src0->ne[2] * src0->nb[2] &&
+            dst->nb[3]  == (size_t) dst->ne[2]  * dst->nb[2]) {
+            struct ggml_tensor src0_c = *src0;
+            struct ggml_tensor dst_c  = *dst;
+            src0_c.ne[2] *= src0_c.ne[3];
+            src0_c.ne[3] = 1;
+            src0_c.nb[3] = (size_t) src0_c.ne[2] * src0_c.nb[2];
+            dst_c.ne[2]  *= dst_c.ne[3];
+            dst_c.ne[3]  = 1;
+            dst_c.nb[3]  = (size_t) dst_c.ne[2] * dst_c.nb[2];
+            node.inputs[0]  = node.add_dummy(src0_c);
+            node.outputs[0] = node.add_dummy(dst_c);
+        }
+        return;
+    }
 }
 
 bool htp_kparams_solve(
@@ -3945,6 +4281,9 @@ bool htp_kparams_solve(
         const size_t src2_size = src2 ? hex_round_up(ggml_nbytes(src2), 128) : 0;
         return htp_kparams_solve_matmul_impl(sess, src0, src1, dst, src2 ? src2->nb[1] : 0, src2_size, node.kparams<htp_mm_kernel_params>());
     }
+    if (node.opcode == HTP_OP_UPSCALE) {
+        return htp_kparams_solve_upscale(sess, src0, dst, node.kparams<htp_upscale_kernel_params>());
+    }
 
     return htp_kparams_solve(sess, node.node, node.kernel_params);
 }
@@ -3955,7 +4294,7 @@ bool htp_kparams_solve(
     void * kparams
 ) {
     const htp_op_code opcode = htp_op_remap(op);
-    if (htp_op_is_unary(opcode) || htp_op_is_glu(opcode)) {
+    if (htp_op_is_unary(opcode) || htp_op_is_glu(opcode) || opcode == HTP_OP_UPSCALE) {
         htp_opnode tmp(opcode, op);
         if (!htp_kparams_solve(sess, tmp)) {
             return false;

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

@@ -29,6 +29,7 @@
 #include "htp/ssm-conv.h"
 #include "htp/sum-rows-ops.h"
 #include "htp/unary-ops.h"
+#include "htp/upscale-ops.h"
 
 #include <type_traits>
 #include <utility>
@@ -134,4 +135,19 @@ bool htp_kparams_solve_unary_mul_glu(
     struct htp_glu_kernel_params * kparams
 );
 
+bool htp_kparams_solve_upscale(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * src0,
+    const struct ggml_tensor * dst,
+    struct htp_upscale_kernel_params * kparams
+);
+
+static inline bool htp_kparams_solve_upscale(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op,
+    struct htp_upscale_kernel_params * kparams
+) {
+    return htp_kparams_solve_upscale(sess, op->src[0], op, kparams);
+}
+
 #endif // HTP_KPARAMS_H

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

@@ -19,6 +19,7 @@
 #include "htp/ssm-conv.h"
 #include "htp/sum-rows-ops.h"
 #include "htp/unary-ops.h"
+#include "htp/upscale-ops.h"
 
 #include <stdio.h>
 
@@ -224,6 +225,19 @@ static std::string format_kernel_params(const htp_opnode & node) {
     } else if (node.opcode == HTP_OP_SOLVE_TRI) {
         const auto * kparams = node.kparams<htp_solve_tri_kernel_params>();
         snprintf(str, max_size, "%s nth %d", kparams->batched ? "batched" : "chunked", (int) kparams->n_threads);
+    } else if (node.opcode == HTP_OP_UPSCALE) {
+        const auto * kparams = node.kparams<htp_upscale_kernel_params>();
+        const char * ktype_str = "unknown";
+        switch (kparams->kernel_type) {
+            case HTP_UPSCALE_KERNEL_NEAREST:     ktype_str = "nearest";     break;
+            case HTP_UPSCALE_KERNEL_NEAREST_2X:  ktype_str = "nearest-2x";  break;
+            case HTP_UPSCALE_KERNEL_NEAREST_4X:  ktype_str = "nearest-4x";  break;
+            case HTP_UPSCALE_KERNEL_BILINEAR:    ktype_str = kparams->align_corners ? "bilinear-ac" : "bilinear"; break;
+            case HTP_UPSCALE_KERNEL_BILINEAR_2X: ktype_str = "bilinear-2x"; break;
+            case HTP_UPSCALE_KERNEL_BILINEAR_AA: ktype_str = "bilinear-aa"; break;
+            case HTP_UPSCALE_KERNEL_BICUBIC:     ktype_str = "bicubic";     break;
+        }
+        snprintf(str, max_size, "%s nth %d chk %d vtcm %u", ktype_str, (int) kparams->n_threads, (int) kparams->dst_nrows_per_chunk, (unsigned int) kparams->vtcm_size);
     } else {
         snprintf(str, max_size, "----");
     }

ggml/src/ggml-hexagon/htp/CMakeLists.txt

@@ -47,6 +47,7 @@ add_library(${HTP_LIB} SHARED
     conv2d-ops.c
     roll-ops.c
     pool-ops.c
+    upscale-ops.c
     allreduce-ops.c
 )
 

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

@@ -47,87 +47,108 @@ static void concat_2d_f32_transposed(unsigned int nth, unsigned int ith, void *
     const uint32_t src0_ne0 = src0->ne[0];
     const uint32_t src1_ne0 = src1->ne[0];
 
-    const uint32_t row_end = cctx->row_start + cctx->nrows;
-    const uint32_t start_i = cctx->row_start + ith * cctx->nrows_per_thread;
-    const uint32_t end_i   = (start_i + cctx->nrows_per_thread < row_end) ? (start_i + cctx->nrows_per_thread) : row_end;
-    if (start_i >= end_i) return;
+    // dst rows [row_begin, row_end) of this thread; a dst row is a src0 row followed by a src1 column
+    const uint32_t row_begin = cctx->row_start + ith * cctx->nrows_per_thread;
+    const uint32_t row_end   = MIN(row_begin + cctx->nrows_per_thread, cctx->row_start + cctx->nrows);
+    if (row_begin >= row_end) return;
 
     dma_queue * dma_q = octx->ctx->dma[ith];
-
-    uint8_t * spad0_base = cctx->spad0_base + ith * cctx->spad0_size_per_thread;
-    uint8_t * spad1_base = cctx->spad1_base + ith * cctx->spad1_size_per_thread;
-
-    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);
-    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);
-
     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];
 
-    uint32_t p = 0;
-    uint32_t i = start_i;
+    // tile: block_rows dst rows x block_cols src1 columns
+    const uint32_t block_rows = 128;
+    const uint32_t block_cols = 256;
+    // +1 element spreads the gather lanes across VTCM, so they do not serialize on one path
+    const uint32_t spad1_stride = block_rows * sizeof(float) + sizeof(float);
 
-    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);
+    const HVX_Vector offsets = hvx_vec_gather_offsets_w(spad1_stride);
+    const uint32_t block_cols_padded = MIN(hex_round_up(src1_ne0, 32), block_cols);
+    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_stride      = src0_row_padded + block_cols_padded * sizeof(float);
 
-    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));
+    // two slots per scratchpad: tile b gathers in slot b&1 while dst[b-1] and src1[b+1] DMA in the other
+    uint8_t * spad0_base = cctx->spad0_base + ith * cctx->spad0_size_per_thread;
+    uint8_t * spad1_base = cctx->spad1_base + ith * cctx->spad1_size_per_thread;
+    const uint32_t spad0_slot = block_rows * spad0_stride;
+    const uint32_t spad1_slot = block_cols_padded * spad1_stride;
 
-    dma_queue_pop(dma_q); // src1
+    // tiles walk (plane, row, col) with col fastest; src0 goes out with the col == 0 tile only
+    uint32_t plane = 0;
+    uint32_t row   = row_begin;
+    uint32_t col   = 0;
+    uint32_t slot  = 0;
+    bool dst_pending = false;
 
-    while (p < cctx->nplanes) {
-        const uint32_t current_block_i = MIN(end_i - i, block_i);
+    const dma_addr_t src1_addr = concat_plane_addr(src1, plane, div_ne2, ne2) + row * src1->nb[1];
+    dma_queue_push(dma_q, dma_make_data(spad1_base, src1_addr), spad1_stride, src1->nb[0], MIN(row_end - row, block_rows) * sizeof(float), MIN(src1_ne0, block_cols));
 
-        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));
+    const dma_addr_t src0_addr = concat_plane_addr(src0, plane, div_ne2, ne2) + row * src0->nb[1];
+    dma_queue_push(dma_q, dma_make_data(spad0_base + src0_pre, src0_addr), spad0_stride, src0->nb[1], src0_row_bytes, MIN(row_end - row, block_rows));
+
+    while (plane < cctx->nplanes) {
+        const uint32_t nrows = MIN(row_end - row, block_rows);
+        const uint32_t ncols = MIN(src1_ne0 - col, block_cols);
+
+        uint8_t * spad0 = spad0_base + slot * spad0_slot;
+        uint8_t * spad1 = spad1_base + slot * spad1_slot;
+
+        dma_queue_pop(dma_q); // src1
+        if (col == 0) {
+            dma_queue_pop(dma_q); // src0
         }
-        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++;
+        uint32_t next_plane = plane;
+        uint32_t next_row   = row;
+        uint32_t next_col   = col + block_cols;
+        if (next_col >= src1_ne0) {
+            next_col = 0;
+            next_row = row + block_rows;
+            if (next_row >= row_end) {
+                next_row = row_begin;
+                next_plane++;
+            }
         }
-        const bool has_next = np < cctx->nplanes;
-        const uint32_t next_block_i = MIN(end_i - ni, block_i);
+        const bool has_next = next_plane < cctx->nplanes;
+        const uint32_t next_nrows = MIN(row_end - next_row, block_rows);
+        const uint32_t next_ncols = MIN(src1_ne0 - next_col, block_cols);
 
-        // spad1 is free after the gather sync, prefetch next src1 ahead of the dst write
+        // other spad1 slot is free since the previous gather sync
         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);
+            const dma_addr_t addr = concat_plane_addr(src1, next_plane, div_ne2, ne2) + next_row * src1->nb[1] + next_col * src1->nb[0];
+            dma_queue_push(dma_q, dma_make_data(spad1_base + (slot ^ 1) * spad1_slot, addr), spad1_stride, src1->nb[0], next_nrows * sizeof(float), next_ncols);
         }
 
-        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
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) row);
+        for (uint32_t c = 0; c < ncols; c += 32) {
+            hvx_transpose_32x32_w_gather(spad0 + src0_row_padded + c * sizeof(float), spad0_stride, spad1 + c * spad1_stride, spad1_stride, offsets, nrows, MIN(ncols - c, 32));
         }
-        dma_queue_pop(dma_q); // dst
+        hvx_gather_sync(spad0 + src0_row_padded);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) row);
 
-        // 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);
+        // other spad0 slot is free once dst[b-1] is written
+        if (dst_pending) {
+            dma_queue_pop(dma_q); // dst
+        }
+        if (has_next && next_col == 0) {
+            const dma_addr_t addr = concat_plane_addr(src0, next_plane, div_ne2, ne2) + next_row * src0->nb[1];
+            dma_queue_push(dma_q, dma_make_data(spad0_base + (slot ^ 1) * spad0_slot + src0_pre, addr), spad0_stride, src0->nb[1], src0_row_bytes, next_nrows);
         }
 
-        p = np;
-        i = ni;
+        // the col == 0 tile also carries the src0 columns in front
+        const uint32_t lead = (col == 0) ? src0_row_bytes : 0;
+        const dma_addr_t dst_addr = concat_plane_addr(dst, plane, div_ne2, ne2) + row * dst->nb[1] + src0_row_bytes + col * sizeof(float) - lead;
+        dma_queue_push(dma_q, dma_make_data(dst_addr, spad0 + src0_row_padded - lead), dst->nb[1], spad0_stride, lead + ncols * sizeof(float), nrows);
+        dst_pending = true;
+
+        plane = next_plane;
+        row   = next_row;
+        col   = next_col;
+        slot ^= 1;
     }
     dma_queue_flush(dma_q);
 }
@@ -143,87 +164,108 @@ static void concat_2d_f16_transposed(unsigned int nth, unsigned int ith, void *
     const uint32_t src0_ne0 = src0->ne[0];
     const uint32_t src1_ne0 = src1->ne[0];
 
-    const uint32_t row_end = cctx->row_start + cctx->nrows;
-    const uint32_t start_i = cctx->row_start + ith * cctx->nrows_per_thread;
-    const uint32_t end_i   = (start_i + cctx->nrows_per_thread < row_end) ? (start_i + cctx->nrows_per_thread) : row_end;
-    if (start_i >= end_i) return;
+    // dst rows [row_begin, row_end) of this thread; a dst row is a src0 row followed by a src1 column
+    const uint32_t row_begin = cctx->row_start + ith * cctx->nrows_per_thread;
+    const uint32_t row_end   = MIN(row_begin + cctx->nrows_per_thread, cctx->row_start + cctx->nrows);
+    if (row_begin >= row_end) return;
 
     dma_queue * dma_q = octx->ctx->dma[ith];
-
-    uint8_t * spad0_base = cctx->spad0_base + ith * cctx->spad0_size_per_thread;
-    uint8_t * spad1_base = cctx->spad1_base + ith * cctx->spad1_size_per_thread;
-
-    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);
-    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);
-
     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];
 
-    uint32_t p = 0;
-    uint32_t i = start_i;
+    // tile: block_rows dst rows x block_cols src1 columns
+    const uint32_t block_rows = 128;
+    const uint32_t block_cols = 512;
+    // +1 element spreads the gather lanes across VTCM, so they do not serialize on one path
+    const uint32_t spad1_stride = block_rows * sizeof(__fp16) + sizeof(__fp16);
 
-    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);
+    const HVX_Vector offsets = hvx_vec_gather_offsets_h(spad1_stride);
+    const uint32_t block_cols_padded = MIN(hex_round_up(src1_ne0, 64), block_cols);
+    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_stride      = src0_row_padded + block_cols_padded * sizeof(__fp16);
 
-    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));
+    // two slots per scratchpad: tile b gathers in slot b&1 while dst[b-1] and src1[b+1] DMA in the other
+    uint8_t * spad0_base = cctx->spad0_base + ith * cctx->spad0_size_per_thread;
+    uint8_t * spad1_base = cctx->spad1_base + ith * cctx->spad1_size_per_thread;
+    const uint32_t spad0_slot = block_rows * spad0_stride;
+    const uint32_t spad1_slot = block_cols_padded * spad1_stride;
 
-    dma_queue_pop(dma_q); // src1
+    // tiles walk (plane, row, col) with col fastest; src0 goes out with the col == 0 tile only
+    uint32_t plane = 0;
+    uint32_t row   = row_begin;
+    uint32_t col   = 0;
+    uint32_t slot  = 0;
+    bool dst_pending = false;
 
-    while (p < cctx->nplanes) {
-        const uint32_t current_block_i = MIN(end_i - i, block_i);
+    const dma_addr_t src1_addr = concat_plane_addr(src1, plane, div_ne2, ne2) + row * src1->nb[1];
+    dma_queue_push(dma_q, dma_make_data(spad1_base, src1_addr), spad1_stride, src1->nb[0], MIN(row_end - row, block_rows) * sizeof(__fp16), MIN(src1_ne0, block_cols));
 
-        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));
+    const dma_addr_t src0_addr = concat_plane_addr(src0, plane, div_ne2, ne2) + row * src0->nb[1];
+    dma_queue_push(dma_q, dma_make_data(spad0_base + src0_pre, src0_addr), spad0_stride, src0->nb[1], src0_row_bytes, MIN(row_end - row, block_rows));
+
+    while (plane < cctx->nplanes) {
+        const uint32_t nrows = MIN(row_end - row, block_rows);
+        const uint32_t ncols = MIN(src1_ne0 - col, block_cols);
+
+        uint8_t * spad0 = spad0_base + slot * spad0_slot;
+        uint8_t * spad1 = spad1_base + slot * spad1_slot;
+
+        dma_queue_pop(dma_q); // src1
+        if (col == 0) {
+            dma_queue_pop(dma_q); // src0
         }
-        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++;
+        uint32_t next_plane = plane;
+        uint32_t next_row   = row;
+        uint32_t next_col   = col + block_cols;
+        if (next_col >= src1_ne0) {
+            next_col = 0;
+            next_row = row + block_rows;
+            if (next_row >= row_end) {
+                next_row = row_begin;
+                next_plane++;
+            }
         }
-        const bool has_next = np < cctx->nplanes;
-        const uint32_t next_block_i = MIN(end_i - ni, block_i);
+        const bool has_next = next_plane < cctx->nplanes;
+        const uint32_t next_nrows = MIN(row_end - next_row, block_rows);
+        const uint32_t next_ncols = MIN(src1_ne0 - next_col, block_cols);
 
-        // spad1 is free after the gather sync, prefetch next src1 ahead of the dst write
+        // other spad1 slot is free since the previous gather sync
         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);
+            const dma_addr_t addr = concat_plane_addr(src1, next_plane, div_ne2, ne2) + next_row * src1->nb[1] + next_col * src1->nb[0];
+            dma_queue_push(dma_q, dma_make_data(spad1_base + (slot ^ 1) * spad1_slot, addr), spad1_stride, src1->nb[0], next_nrows * sizeof(__fp16), next_ncols);
         }
 
-        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
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) row);
+        for (uint32_t c = 0; c < ncols; c += 64) {
+            hvx_transpose_64x64_h_gather(spad0 + src0_row_padded + c * sizeof(__fp16), spad0_stride, spad1 + c * spad1_stride, spad1_stride, offsets, nrows, MIN(ncols - c, 64));
         }
-        dma_queue_pop(dma_q); // dst
+        hvx_gather_sync(spad0 + src0_row_padded);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) row);
 
-        // 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);
+        // other spad0 slot is free once dst[b-1] is written
+        if (dst_pending) {
+            dma_queue_pop(dma_q); // dst
+        }
+        if (has_next && next_col == 0) {
+            const dma_addr_t addr = concat_plane_addr(src0, next_plane, div_ne2, ne2) + next_row * src0->nb[1];
+            dma_queue_push(dma_q, dma_make_data(spad0_base + (slot ^ 1) * spad0_slot + src0_pre, addr), spad0_stride, src0->nb[1], src0_row_bytes, next_nrows);
         }
 
-        p = np;
-        i = ni;
+        // the col == 0 tile also carries the src0 columns in front
+        const uint32_t lead = (col == 0) ? src0_row_bytes : 0;
+        const dma_addr_t dst_addr = concat_plane_addr(dst, plane, div_ne2, ne2) + row * dst->nb[1] + src0_row_bytes + col * sizeof(__fp16) - lead;
+        dma_queue_push(dma_q, dma_make_data(dst_addr, spad0 + src0_row_padded - lead), dst->nb[1], spad0_stride, lead + ncols * sizeof(__fp16), nrows);
+        dst_pending = true;
+
+        plane = next_plane;
+        row   = next_row;
+        col   = next_col;
+        slot ^= 1;
     }
     dma_queue_flush(dma_q);
 }

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

@@ -53,13 +53,17 @@ static inline void htp_concat_transposed_vtcm_layout_build(
     uint32_t type_size,
     uint32_t n_threads) {
 
-    uint32_t block_i = (type_size == 4) ? 32 : 64;
-    uint32_t spad1_stride = block_i * type_size;
-    uint32_t src1_ne0_padded = hex_round_up(src1_ne0, block_i);
-    uint32_t spad0_row_bytes = hex_round_up(src0_ne0 * type_size, 128) + src1_ne0_padded * type_size;
+    // src1_ne0 is processed in chunks of block_cols tokens, so VTCM does not grow with the batch
+    uint32_t block_rows = 128;
+    uint32_t block_cols = (type_size == 4) ? 256 : 512;
+    // +1 element spreads the gather lanes across VTCM, so they do not serialize on one path
+    uint32_t spad1_stride = block_rows * type_size + type_size;
+    uint32_t block_cols_padded = MIN(hex_round_up(src1_ne0, (type_size == 4) ? 32 : 64), block_cols);
+    uint32_t spad0_row_bytes = hex_round_up(src0_ne0 * type_size, 128) + block_cols_padded * type_size;
 
-    L->src0_spad_size_per_thread = block_i * spad0_row_bytes;
-    L->src1_spad_size_per_thread = src1_ne0_padded * spad1_stride;
+    // double-buffered: two slots per thread
+    L->src0_spad_size_per_thread = 2 * block_rows * spad0_row_bytes;
+    L->src1_spad_size_per_thread = 2 * block_cols_padded * spad1_stride;
 
     size_t off = 0;
     VTCM_LAYOUT_ALLOC(off, off_spad0, (size_t) n_threads * L->src0_spad_size_per_thread);

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

@@ -35,6 +35,10 @@ struct htp_copy_context {
     struct htp_copy_convert_vtcm_layout layout;
     uint8_t * vtcm_src0;
     uint8_t * vtcm_dst;
+
+    uint32_t tile_start;
+    uint32_t ntiles;
+    uint32_t tiles_per_thread;
 };
 
 #define cpy_preamble                              \
@@ -162,6 +166,171 @@ static inline void cpy_cursor_advance(struct cpy_row_cursor * c, uint32_t n, uin
     }
 }
 
+struct cpy_tile_cursor {
+    uint32_t ibk;
+    uint32_t ib0;
+    uint32_t i2;
+    uint32_t i3;
+};
+
+static inline void cpy_tile_cursor_advance(struct cpy_tile_cursor * c, const struct htp_copy_transpose_params * tp) {
+    if (++c->ibk == tp->nkb) {
+        c->ibk = 0;
+        if (++c->ib0 == tp->n0b) {
+            c->ib0 = 0;
+            if (++c->i2 == tp->ne[2]) {
+                c->i2 = 0;
+                c->i3++;
+            }
+        }
+    }
+}
+
+// Gather offsets for element e = l >> log2_m at (e % r, e / r)
+static inline HVX_Vector cpy_transpose_offsets(const struct htp_copy_transpose_params * tp) {
+    const uint32_t r_mask = (1u << tp->log2_r) - 1;
+    if (tp->elem_size != 2) {
+        const uint32_t m_mask = (1u << tp->log2_m) - 1;
+        int32_t offsets[32] __attribute__((aligned(128)));
+        for (uint32_t l = 0; l < 32; l++) {
+            const uint32_t e = l >> tp->log2_m;
+            offsets[l] = (e & r_mask) * tp->in_stride + (e >> tp->log2_r) * tp->elem_size + (l & m_mask) * 4;
+        }
+        return *(HVX_Vector *) offsets;
+    }
+    int16_t offsets[64] __attribute__((aligned(128)));
+    for (uint32_t l = 0; l < 64; l++) {
+        offsets[l] = (l & r_mask) * tp->in_stride + (l >> tp->log2_r) * 2;
+    }
+    return *(HVX_Vector *) offsets;
+}
+
+static inline void cpy_transpose_push_in(dma_queue * q, const struct htp_copy_transpose_params * tp, dma_addr_t src,
+                                         const struct cpy_tile_cursor * c, uint8_t * buf) {
+    const uint32_t a0  = c->ib0 * tp->blk0;
+    const uint32_t ak  = c->ibk * tp->blkk;
+    const uint32_t off = a0 * tp->src_nb[0] + ak * tp->elem_size + c->i2 * tp->src_nb[1] + c->i3 * tp->src_nb[2];
+    dma_queue_push(q, dma_make_data(buf, src + off), tp->in_stride, tp->src_nb[0],
+                   MIN(tp->blkk, tp->ne[1] - ak) * tp->elem_size, MIN(tp->blk0, tp->ne[0] - a0));
+}
+
+// Pipelined tile worker: DMA in -> vgather transpose -> DMA out
+static void cpy_thread_transpose(unsigned int nth, unsigned int ith, void * data) {
+    struct htp_copy_context * ct = (struct htp_copy_context *) data;
+    struct htp_ops_context * octx = ct->octx;
+    const struct htp_copy_transpose_params * tp = &ct->kparams->u.transpose;
+
+    const uint32_t t0 = ct->tile_start + ith * ct->tiles_per_thread;
+    const uint32_t t1 = MIN(t0 + ct->tiles_per_thread, ct->tile_start + ct->ntiles);
+    if (t0 >= t1) {
+        return;
+    }
+
+    dma_queue * dma_q = octx->ctx->dma[ith];
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
+    const dma_addr_t src = octx->src[0]->data;
+    const dma_addr_t dst = octx->dst->data;
+
+    const uint32_t es         = tp->elem_size;
+    const uint32_t ne0        = tp->ne[0];
+    const uint32_t nek        = tp->ne[1];
+    const uint32_t in_stride  = tp->in_stride;
+    const uint32_t out_stride = tp->out_stride;
+
+    const HVX_Vector offsets   = cpy_transpose_offsets(tp);
+    const uint32_t packed_mu   = (ne0 - 1) * in_stride + 127;
+    const uint32_t packed_step = 128 >> tp->log2_r;
+    const uint32_t epv         = 1u << tp->log2_r;
+
+    uint8_t * in_buf[2];
+    uint8_t * out_buf[2];
+    in_buf[0]  = ct->vtcm_src0 + ith * 2 * tp->in_buf_size;
+    in_buf[1]  = in_buf[0] + tp->in_buf_size;
+    out_buf[0] = ct->vtcm_dst + ith * 2 * tp->out_buf_size;
+    out_buf[1] = out_buf[0] + tp->out_buf_size;
+
+    struct cpy_tile_cursor cur;
+    {
+        const uint32_t q0 = fastdiv(t0, &tp->div_nkb);
+        const uint32_t q1 = fastdiv(q0, &tp->div_n0b);
+        cur.ibk = t0 - q0 * tp->nkb;
+        cur.ib0 = q0 - q1 * tp->n0b;
+        cur.i3  = fastdiv(q1, &tp->div_ne2);
+        cur.i2  = q1 - cur.i3 * tp->ne[2];
+    }
+
+    struct cpy_tile_cursor nxt = cur;
+
+    cpy_transpose_push_in(dma_q, tp, src, &cur, in_buf[0]);
+
+    for (uint32_t t = t0, i = 0; t < t1; t++, i++) {
+        while (dma_queue_depth(dma_q) > (i > 0 ? 1u : 0u)) {
+            dma_queue_pop(dma_q);
+        }
+
+        uint8_t * in  = in_buf[i & 1];
+        uint8_t * out = out_buf[i & 1];
+
+        const uint32_t a0 = cur.ib0 * tp->blk0;
+        const uint32_t ak = cur.ibk * tp->blkk;
+        const uint32_t n0 = MIN(tp->blk0, ne0 - a0);
+        const uint32_t nk = MIN(tp->blkk, nek - ak);
+
+        if (t + 1 < t1) {
+            cpy_tile_cursor_advance(&nxt, tp);
+            cpy_transpose_push_in(dma_q, tp, src, &nxt, in_buf[(i + 1) & 1]);
+        }
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) t);
+        if (tp->packed) {
+            const uint32_t nvec = (nk * ne0 * es + 127) >> 7;
+            if (es != 2) {
+                #pragma unroll(4)
+                for (uint32_t v = 0; v < nvec; v++) {
+                    Q6_vgather_ARMVw((HVX_Vector *) (out + v * 128), (size_t) (in + v * packed_step), packed_mu, offsets);
+                }
+            } else {
+                #pragma unroll(4)
+                for (uint32_t v = 0; v < nvec; v++) {
+                    Q6_vgather_ARMVh((HVX_Vector *) (out + v * 128), (size_t) (in + v * packed_step), packed_mu, offsets);
+                }
+            }
+        } else {
+            for (uint32_t rb = 0; rb < n0; rb += epv) {
+                const uint32_t  mu = (MIN(epv, n0 - rb) - 1) * in_stride + es - 1;
+                uint8_t *       d  = out + rb * es;
+                const uint8_t * s  = in + rb * in_stride;
+                if (es != 2) {
+                    #pragma unroll(4)
+                    for (uint32_t ik = 0; ik < nk; ik++) {
+                        Q6_vgather_ARMVw((HVX_Vector *) (d + ik * out_stride), (size_t) (s + ik * es), mu, offsets);
+                    }
+                } else {
+                    #pragma unroll(4)
+                    for (uint32_t ik = 0; ik < nk; ik++) {
+                        Q6_vgather_ARMVh((HVX_Vector *) (d + ik * out_stride), (size_t) (s + ik * 2), mu, offsets);
+                    }
+                }
+            }
+        }
+        hvx_gather_sync(out);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) t);
+
+        const uint32_t doff = a0 * es + ak * tp->dst_nb[0] + cur.i2 * tp->dst_nb[1] + cur.i3 * tp->dst_nb[2];
+        if (tp->packed) {
+            const uint32_t n = nk * ne0 * es;
+            dma_queue_push(dma_q, dma_make_data(dst + doff, out), n, n, n, 1);
+        } else {
+            dma_queue_push(dma_q, dma_make_data(dst + doff, out), tp->dst_nb[0], out_stride, n0 * es, nk);
+        }
+
+        cur = nxt;
+    }
+
+    dma_queue_flush(dma_q);
+}
+
 #define DEFINE_CPY_CONVERT_SAMESHAPE(NAME, CONV_FUNC)                                        \
 static void cpy_thread_##NAME##_sameshape(unsigned int nth, unsigned int ith, void * data) { \
     struct htp_copy_context * ct = (struct htp_copy_context *) data;                         \
@@ -439,6 +608,48 @@ static int cpy_reshape(struct htp_ops_context * octx, const struct htp_copy_kern
     return HTP_STATUS_OK;
 }
 
+static int cpy_transpose(struct htp_ops_context * octx, const struct htp_copy_kernel_params * kparams) {
+    if (!htp_ops_context_set_n_threads(octx, kparams->n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    const struct htp_copy_transpose_params * tp = &kparams->u.transpose;
+
+    uint32_t tile_start = 0;
+    uint32_t ntiles     = tp->total_tiles;
+
+    if (octx->ctx->mdev.count > 1) {
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(
+            tp->total_tiles, 1, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        tile_start = range.start;
+        ntiles     = range.count;
+    }
+
+    if (ntiles == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    if (kparams->vtcm_size > octx->ctx->vtcm_size) {
+        return HTP_STATUS_VTCM_TOO_SMALL;
+    }
+
+    const uint32_t n_threads = octx->n_threads;
+
+    struct htp_copy_context ct;
+    ct.octx             = octx;
+    ct.kparams          = kparams;
+    ct.tile_start       = tile_start;
+    ct.ntiles           = ntiles;
+    ct.tiles_per_thread = fastdiv(ntiles + n_threads - 1, &octx->n_threads_div);
+
+    uint8_t * vtcm_base = (uint8_t *) octx->ctx->vtcm_base;
+    ct.vtcm_src0 = vtcm_base;
+    ct.vtcm_dst  = vtcm_base + (size_t) n_threads * 2 * tp->in_buf_size;
+
+    work_queue_run(octx->ctx->work_queue, cpy_thread_transpose, &ct, n_threads);
+    return HTP_STATUS_OK;
+}
+
 int op_cpy(struct htp_ops_context * octx) {
     const struct htp_copy_kernel_params * kparams = (const struct htp_copy_kernel_params *) octx->kernel_params;
     int status = HTP_STATUS_OK;
@@ -459,6 +670,9 @@ int op_cpy(struct htp_ops_context * octx) {
         case HTP_COPY_KERNEL_RESHAPE:
             status = cpy_reshape(octx, kparams);
             break;
+        case HTP_COPY_KERNEL_TRANSPOSE:
+            status = cpy_transpose(octx, kparams);
+            break;
         case HTP_COPY_KERNEL_GATHER_ROWS:
             status = htp_gather_rows(octx, octx->src[0], NULL, &kparams->u.gather, kparams->n_threads);
             break;

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

@@ -15,6 +15,7 @@ enum htp_copy_kernel_type {
     HTP_COPY_KERNEL_RESHAPE            = 4,
     HTP_COPY_KERNEL_SCALAR             = 5,
     HTP_COPY_KERNEL_GATHER_ROWS        = 6,
+    HTP_COPY_KERNEL_TRANSPOSE          = 7,
 };
 
 struct htp_copy_convert_params {
@@ -32,6 +33,29 @@ struct htp_copy_reshape_params {
     struct fastdiv_values div_ne02_ne01_ne00;
 };
 
+// Batched 2D transpose of elem_size elements (dim 0 dense in dst, dim 1 dense in src)
+struct htp_copy_transpose_params {
+    uint32_t ne[3];          // dims 0..2
+    uint32_t src_nb[3];      // src strides for dims 0, 2, 3
+    uint32_t dst_nb[3];      // dst strides for dims 1, 2, 3
+    struct fastdiv_values div_nkb;
+    struct fastdiv_values div_n0b;
+    struct fastdiv_values div_ne2;
+    uint32_t blk0;           // dim-0 elements per tile
+    uint32_t blkk;           // dim-1 elements per tile
+    uint32_t n0b;            // tiles along dim 0
+    uint32_t nkb;            // tiles along dim 1
+    uint32_t in_stride;      // VTCM row pitch of staged src tile
+    uint32_t out_stride;     // VTCM row pitch of transposed tile
+    uint32_t in_buf_size;
+    uint32_t out_buf_size;
+    uint32_t total_tiles;
+    uint8_t  elem_size;
+    uint8_t  packed;         // gather straight to dst order
+    uint8_t  log2_m;         // log2(gather lanes per element)
+    uint8_t  log2_r;         // log2(dim-0 elements per vector)
+};
+
 struct htp_copy_kernel_params {
     uint8_t  kernel_type;
     uint8_t  src0_type_size;
@@ -46,6 +70,7 @@ struct htp_copy_kernel_params {
         struct htp_copy_convert_params convert;
         struct htp_copy_reshape_params reshape;
         struct htp_gather_rows_params  gather;
+        struct htp_copy_transpose_params transpose;
     } u;
 };
 
@@ -116,6 +141,76 @@ static inline void htp_copy_convert_solve_layout(
     htp_copy_convert_vtcm_layout_build(layout, ne00, src_type_size, dst_type_size, blk_rows, n_threads);
 }
 
+#define HTP_COPY_TRANSPOSE_BLK_BYTES (64 * 1024)
+
+static inline bool htp_copy_transpose_layout_build(struct htp_copy_transpose_params * tp, uint32_t blk_bytes) {
+    const uint32_t ne0  = tp->ne[0];
+    const uint32_t nek  = tp->ne[1];
+    const uint32_t es   = tp->elem_size;
+    const uint32_t epv  = 128 / es;
+
+    uint32_t blk0;
+    uint32_t blkk;
+    if (tp->packed) {
+        // Packed: vector holds epv / ne0 whole dst rows
+        blk0 = ne0;
+        blkk = blk_bytes / (ne0 * es);
+        if (es == 2 && blkk > 16384 / ne0 - 64) {
+            blkk = 16384 / ne0 - 64;
+        }
+        if (blkk > nek) {
+            blkk = nek;
+        }
+        if (blkk == 0) {
+            return false;
+        }
+    } else if (nek <= epv) {
+        // Stage short contiguous dim and expand dim 0
+        blkk = nek;
+        blk0 = blk_bytes / (nek * es);
+        blk0 -= blk0 % epv;
+        if (blk0 > hex_round_up(ne0, epv)) {
+            blk0 = hex_round_up(ne0, epv);
+        }
+        if (blk0 < epv) {
+            return false;
+        }
+    } else {
+        blk0 = hex_round_up(ne0, epv);
+        if (blk0 > 256) {
+            blk0 = 256;
+        }
+        blkk = blk_bytes / (blk0 * es);
+        blkk -= blkk % epv;
+        // Cap f16 blkk to 256 for 16-bit gather offsets
+        const uint32_t max_blkk = (es == 2 ? 512 : 4096) / es;
+        if (blkk > max_blkk) {
+            blkk = max_blkk;
+        }
+        if (blkk > hex_round_up(nek, epv)) {
+            blkk = hex_round_up(nek, epv);
+        }
+        if (blkk < epv) {
+            return false;
+        }
+    }
+
+    tp->blk0         = blk0;
+    tp->blkk         = blkk;
+    tp->n0b          = (ne0 + blk0 - 1) / blk0;
+    tp->nkb          = (nek + blkk - 1) / blkk;
+    // Pad row pitch by 1 element to avoid VTCM bank conflicts
+    const uint32_t pad = (es == 2) ? 2 : 4;
+    tp->in_stride    = blkk * es + pad;
+    tp->out_stride   = blk0 * es;
+    // Extra vector for packed gather overhang
+    tp->in_buf_size  = hex_round_up(blk0 * tp->in_stride + (tp->packed ? 128 : 0), 128);
+    tp->out_buf_size = hex_round_up(blkk * tp->out_stride, 128);
+    tp->log2_m       = es == 2 ? 0 : __builtin_ctz(es) - 2;
+    tp->log2_r       = __builtin_ctz(tp->packed ? ne0 : epv);
+    return true;
+}
+
 #if defined(__cplusplus)
 static_assert(sizeof(struct htp_copy_kernel_params) <= 128, "htp_copy_kernel_params is too large for kernel_params blob");
 #else

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

@@ -141,43 +141,9 @@ bool dma_queue_push_fallback_2d(dma_queue * q, dma_data ddata, size_t dst_stride
     dma_ring_flush(r1);
     r0->tail = r1->tail;
 
-    return dma_ring_push_single_2d(r0, ddata, 0, 0, 0, /*nrows=*/ 0);
-}
-
-bool dma_queue_push_fallback_contig(dma_queue * q, dma_data ddata, size_t total) {
-    dma_ring * r0 = q->ring0;
-    dma_ring * r1 = q->ring1;
-
-    if (((r0->push_idx + 1) & r0->idx_mask) == r0->pop_idx) {
-        return false;
-    }
-
-    r1->tail = r0->tail;
-
-    size_t rem_bytes   = total;
-    dma_addr_t cur_dst = ddata.dst;
-    dma_addr_t cur_src = ddata.src;
-
-    while (rem_bytes > 0) {
-        const uint32_t cur_bytes = MIN(rem_bytes, DMA_SAFE_CHUNK_SIZE);
-        dma_data cur_data = dma_make_data(cur_dst, cur_src);
-        if (!dma_ring_push_single_contig(r1, cur_data, cur_bytes)) {
-            dma_ring_flush(r1);
-            dma_ring_push_single_contig(r1, cur_data, cur_bytes);
-        }
-        cur_dst   += cur_bytes;
-        cur_src   += cur_bytes;
-        rem_bytes -= cur_bytes;
-    }
-
-    dma_ring_flush(r1);
-    r0->tail = r1->tail;
-
     return dma_ring_push_single_1d(r0, ddata, /*size=*/ 0);
 }
 
-#if __HVX_ARCH__ < 75
-
 bool dma_queue_push_fallback_1d(dma_queue * q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows) {
     dma_ring * r0 = q->ring0;
     dma_ring * r1 = q->ring1;
@@ -188,19 +154,28 @@ bool dma_queue_push_fallback_1d(dma_queue * q, dma_data ddata, size_t dst_stride
 
     r1->tail = r0->tail;
 
-    size_t rem_rows    = nrows;
-    dma_addr_t cur_dst = ddata.dst;
-    dma_addr_t cur_src = ddata.src;
+    dma_addr_t row_dst = ddata.dst;
+    dma_addr_t row_src = ddata.src;
 
-    while (rem_rows > 0) {
-        dma_data cur_data = dma_make_data(cur_dst, cur_src);
-        if (!dma_ring_push_single_1d(r1, cur_data, row_size)) {
-            dma_ring_flush(r1);
-            dma_ring_push_single_1d(r1, cur_data, row_size);
+    for (size_t i = 0; i < nrows; i++) {
+        size_t     rem_bytes = row_size;
+        dma_addr_t cur_dst   = row_dst;
+        dma_addr_t cur_src   = row_src;
+
+        while (rem_bytes > 0) {
+            const uint32_t cur_bytes = MIN(rem_bytes, DMA_SAFE_CHUNK_SIZE);
+            dma_data cur_data = dma_make_data(cur_dst, cur_src);
+            if (!dma_ring_push_single_contig(r1, cur_data, cur_bytes)) {
+                dma_ring_flush(r1);
+                dma_ring_push_single_contig(r1, cur_data, cur_bytes);
+            }
+            cur_dst   += cur_bytes;
+            cur_src   += cur_bytes;
+            rem_bytes -= cur_bytes;
         }
-        cur_dst  += dst_stride;
-        cur_src  += src_stride;
-        rem_rows -= 1;
+
+        row_dst += dst_stride;
+        row_src += src_stride;
     }
 
     dma_ring_flush(r1);
@@ -208,5 +183,3 @@ bool dma_queue_push_fallback_1d(dma_queue * q, dma_data ddata, size_t dst_stride
 
     return dma_ring_push_single_1d(r0, ddata, /*size=*/ 0);
 }
-
-#endif

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

@@ -148,10 +148,7 @@ dma_queue_t dma_queue_alias_init(void * ptr, dma_queue_t main_q);
 void        dma_queue_alias_free(dma_queue_t q);
 
 bool        dma_queue_push_fallback_2d(dma_queue * q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows);
-bool        dma_queue_push_fallback_contig(dma_queue * q, dma_data ddata, size_t total);
-#if __HVX_ARCH__ < 75
 bool        dma_queue_push_fallback_1d(dma_queue * q, dma_data ddata, size_t dst_stride, size_t src_stride, size_t row_size, size_t nrows);
-#endif
 
 // TODO: technically we don't need these and could use Q6_dmstart/wait/etc instead
 // but those do not seem to always compiler properly.
@@ -400,7 +397,7 @@ static inline bool dma_queue_push(dma_queue * q, dma_data ddata, size_t dst_stri
         if (__builtin_expect(total <= DMA_MAX_SIZE_24B, 1)) {
             return dma_ring_push_single_contig(q->ring0, ddata, total);
         }
-        return dma_queue_push_fallback_contig(q, ddata, total);
+        return dma_queue_push_fallback_1d(q, ddata, 0, 0, total, 1);
     }
 
     // 2. Hot path: Standard strided 2D (10-20% of calls)
@@ -410,14 +407,10 @@ static inline bool dma_queue_push(dma_queue * q, dma_data ddata, size_t dst_stri
     }
 
     // 3. Cold path: Descriptor chunking fallbacks (< 0.1%)
-#if __HVX_ARCH__ < 75
-    if (row_size <= DMA_MAX_SIZE_16B && (src_stride | dst_stride) <= DMA_MAX_STRIDE_16B) {
+    if ((row_size | src_stride | dst_stride) <= DMA_MAX_2D_ROW_SIZE) {
         return dma_queue_push_fallback_2d(q, ddata, dst_stride, src_stride, row_size, nrows);
     }
     return dma_queue_push_fallback_1d(q, ddata, dst_stride, src_stride, row_size, nrows);
-#else
-    return dma_queue_push_fallback_2d(q, ddata, dst_stride, src_stride, row_size, nrows);
-#endif
 }
 
 static inline void dma_sync_read(dma_queue * dma_q, void * dst, dma_addr_t src, size_t bytes) {

ggml/src/ggml-hexagon/htp/gated-delta-net-ops.c

@@ -1456,9 +1456,10 @@ static inline void gdn_dma_push_chunk_gb(
     const dma_addr_t g_dma    = g->data + (uint64_t) iv3 * g->nb[3] + (uint64_t) t_chunk * g->nb[2] + (uint64_t) iv1 * g->nb[1];
     const dma_addr_t beta_dma = beta->data + (uint64_t) iv3 * beta->nb[3] + (uint64_t) t_chunk * beta->nb[2] + (uint64_t) iv1 * beta->nb[1];
     const uint32_t row_bytes  = n_batch * sizeof(float);
+    const uint32_t dst_stride = (n_batch > 1) ? row_bytes + sizeof(float) : row_bytes;
 
-    dma_queue_push(dma_q, dma_make_data(vtcm_g_raw, g_dma), row_bytes, g->nb[2], row_bytes, chunk_size);
-    dma_queue_push(dma_q, dma_make_data(vtcm_b_raw, beta_dma), row_bytes, beta->nb[2], row_bytes, chunk_size);
+    dma_queue_push(dma_q, dma_make_data(vtcm_g_raw, g_dma), dst_stride, g->nb[2], row_bytes, chunk_size);
+    dma_queue_push(dma_q, dma_make_data(vtcm_b_raw, beta_dma), dst_stride, beta->nb[2], row_bytes, chunk_size);
 }
 
 static inline void gdn_pack_s_col_tiles(
@@ -1707,19 +1708,16 @@ static void gdn_hvx_phase1a_worker(unsigned int n, unsigned int i, void * data)
         hvx_vmem(head->b_f32[curr_buf] + 0)  = hvx_vmem(bctx->vtcm_b_raw + 0);
         hvx_vmem(head->b_f32[curr_buf] + 32) = hvx_vmem(bctx->vtcm_b_raw + 32);
     } else {
-        int32_t offsets[32] __attribute__((aligned(128)));
-        for (int k = 0; k < 32; ++k) {
-            offsets[k] = k * n_batch * sizeof(float);
-        }
-        HVX_Vector vv = *(const HVX_Vector *) offsets;
+        const uint32_t vtcm_stride = (n_batch + 1) * sizeof(float);
+        const HVX_Vector vv = hvx_vec_gather_offsets_w(vtcm_stride);
         const size_t rt_g = (size_t) ((const uint8_t *) bctx->vtcm_g_raw + i * sizeof(float));
         const size_t rt_b = (size_t) ((const uint8_t *) bctx->vtcm_b_raw + i * sizeof(float));
-        const size_t mu   = 64 * n_batch * sizeof(float);
+        const size_t mu   = 64 * vtcm_stride;
 
         Q6_vgather_ARMVw((HVX_Vector *) (head->g_f32[curr_buf] + 0),  rt_g, mu, vv);
-        Q6_vgather_ARMVw((HVX_Vector *) (head->g_f32[curr_buf] + 32), rt_g + 32 * n_batch * sizeof(float), mu, vv);
+        Q6_vgather_ARMVw((HVX_Vector *) (head->g_f32[curr_buf] + 32), rt_g + 32 * vtcm_stride, mu, vv);
         Q6_vgather_ARMVw((HVX_Vector *) (head->b_f32[curr_buf] + 0),  rt_b, mu, vv);
-        Q6_vgather_ARMVw((HVX_Vector *) (head->b_f32[curr_buf] + 32), rt_b + 32 * n_batch * sizeof(float), mu, vv);
+        Q6_vgather_ARMVw((HVX_Vector *) (head->b_f32[curr_buf] + 32), rt_b + 32 * vtcm_stride, mu, vv);
     }
 
     const uint32_t t_chunk = bctx->c * 64;

ggml/src/ggml-hexagon/htp/gated-delta-net-ops.h

@@ -204,7 +204,7 @@ static inline void htp_gdn_hmx_vtcm_layout_build(
     VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_g_f32[1], bh * dma_scalar_sz, pipeline);
     VTCM_LAYOUT_ALLOC(off, off_b_f32[0], bh * dma_scalar_sz);
     VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_b_f32[1], bh * dma_scalar_sz, pipeline);
-    const size_t raw_gb_sz = hex_round_up(bh * chunk_size * sizeof(float), 128);
+    const size_t raw_gb_sz = hex_round_up((bh + 1) * chunk_size * sizeof(float), 128);
     VTCM_LAYOUT_ALLOC(off, off_g_raw[0], raw_gb_sz);
     VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_g_raw[1], raw_gb_sz, pipeline);
     VTCM_LAYOUT_ALLOC(off, off_b_raw[0], raw_gb_sz);

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

@@ -188,5 +188,6 @@ int op_pool_2d(struct htp_ops_context * octx);
 int op_pool_1d(struct htp_ops_context * octx);
 int op_allreduce(struct htp_ops_context * octx);
 int op_roll(struct htp_ops_context * octx);
+int op_upscale(struct htp_ops_context * octx);
 
 #endif /* HTP_CTX_H */

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

@@ -116,6 +116,7 @@ enum htp_op_code {
     HTP_OP_GLU_SIGMOID,
     HTP_OP_UNARY_CEIL,
     HTP_OP_UNARY_SGN,
+    HTP_OP_UPSCALE,
 
     HTP_OP_INVALID
 };

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

@@ -1049,6 +1049,9 @@ static int execute_op(struct htp_ops_context * octx) {
         case HTP_OP_TRI:
             return op_unary(octx);
 
+        case HTP_OP_UPSCALE:
+            return op_upscale(octx);
+
         case HTP_OP_INVALID:
             break;
     }

ggml/src/ggml-hexagon/htp/ssm-conv.c

@@ -173,11 +173,12 @@ static inline void hvx_ssm_conv_decode_4(const float * x, const float * w, float
 // the lanes past the channel tail.
 static inline void hvx_ssm_conv_transpose_block(const float * raw_block,
                                                 float *       T,
+                                                uint32_t      raw_stride_bytes,
                                                 uint32_t      ncs,
                                                 uint32_t      cb_n,
                                                 HVX_Vector    vv) {
     const size_t   base = (size_t) raw_block;
-    const uint32_t mu   = cb_n * ncs * sizeof(float) - 1;
+    const uint32_t mu   = cb_n * raw_stride_bytes - 1;
 
     for (uint32_t t = 0; t < ncs; ++t) {
         Q6_vgather_ARMVw((HVX_Vector *) (T + (size_t) t * VLEN_FP32), base + t * sizeof(float), mu, vv);
@@ -393,7 +394,8 @@ static void ssm_conv_thread_f32_prefill(unsigned int nth, unsigned int ith, void
 
     // src0 spad holds two raw tiles (fetch of tile n+1 overlaps compute of tile n) plus
     // the transposed block. dst spad holds two tiles so a writeback can stay in flight.
-    const size_t src0_tile_raw_bytes = hex_round_up(d_inner_tile * ncs * sizeof(float), 128);
+    const size_t raw_stride_bytes    = ncs * sizeof(float) + sizeof(float);
+    const size_t src0_tile_raw_bytes = hex_round_up(d_inner_tile * raw_stride_bytes, 128);
     const size_t dst_tile_bytes      = hex_round_up(d_inner_tile * n_t * sizeof(float), 128);
 
     float * src0_tile_raw[2] = { (float *) src0_spad_base, (float *) (src0_spad_base + src0_tile_raw_bytes) };
@@ -419,7 +421,7 @@ static void ssm_conv_thread_f32_prefill(unsigned int nth, unsigned int ith, void
                                           (ir0 + tile_off_fetch) * src0_stride_inner_bytes;       \
             dma_queue_push(dma_q,                                                                 \
                            dma_make_data((uint8_t *) src0_tile_raw[(c) & 1], fetch_ddr),          \
-                           row_bytes, src0_stride_inner_bytes, row_bytes, cur_tile_n);            \
+                           raw_stride_bytes, src0_stride_inner_bytes, row_bytes, cur_tile_n);     \
             tile_off_fetch += d_inner_tile;                                                       \
             if (tile_off_fetch >= d_inner_per_thread) {                                           \
                 tile_off_fetch = 0;                                                               \
@@ -446,11 +448,7 @@ static void ssm_conv_thread_f32_prefill(unsigned int nth, unsigned int ith, void
     const uint32_t C_TILE = VLEN_FP32;
 
     // gather offsets: lane r reads channel r of the raw tile block
-    uint32_t __attribute__((aligned(VLEN))) gather_off[VLEN_FP32];
-    for (uint32_t r = 0; r < VLEN_FP32; ++r) {
-        gather_off[r] = r * ncs * sizeof(float);
-    }
-    const HVX_Vector vv = *(const HVX_Vector *) gather_off;
+    const HVX_Vector vv = hvx_vec_gather_offsets_w(raw_stride_bytes);
 
     uint32_t i3       = 0;
     uint32_t tile_off = 0;
@@ -474,7 +472,8 @@ static void ssm_conv_thread_f32_prefill(unsigned int nth, unsigned int ith, void
             const uint32_t cb_n = MIN(C_TILE, tile_n - cb);
 
             htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_PREP, (uint16_t) tile_off);
-            hvx_ssm_conv_transpose_block(raw + (size_t) cb * ncs, src0_T, ncs, cb_n, vv);
+            const float * raw_ch = (const float *) ((const uint8_t *) raw + (size_t) cb * raw_stride_bytes);
+            hvx_ssm_conv_transpose_block(raw_ch, src0_T, raw_stride_bytes, ncs, cb_n, vv);
             htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_PREP, (uint16_t) tile_off);
 
             const float * restrict wp = src1_T + tile_off + cb;

ggml/src/ggml-hexagon/htp/ssm-conv.h

@@ -54,11 +54,12 @@ static inline void htp_ssm_conv_vtcm_layout_build(
         L->src0_bytes_per_thread         = 2 * src0_tile_raw_bytes + src0_T_bytes;
         L->dst_bytes_per_thread          = 2 * hex_round_up((size_t) d_inner_tile * sizeof(float), 128);
     } else {
-        const uint32_t ncs          = d_conv - 1 + n_t;
-        const size_t   src0_block_T = hex_round_up((size_t) ncs * 32 * sizeof(float), 128);
-        const size_t   src0_tile_raw = hex_round_up((size_t) d_inner_tile * ncs * sizeof(float), 128);
-        L->src0_bytes_per_thread    = 2 * src0_tile_raw + src0_block_T;
-        L->dst_bytes_per_thread     = 2 * hex_round_up((size_t) d_inner_tile * n_t * sizeof(float), 128);
+        const uint32_t ncs              = d_conv - 1 + n_t;
+        const size_t   raw_stride_bytes = ncs * sizeof(float) + sizeof(float);
+        const size_t   src0_block_T     = hex_round_up((size_t) ncs * 32 * sizeof(float), 128);
+        const size_t   src0_tile_raw    = hex_round_up((size_t) d_inner_tile * raw_stride_bytes, 128);
+        L->src0_bytes_per_thread        = 2 * src0_tile_raw + src0_block_T;
+        L->dst_bytes_per_thread         = 2 * hex_round_up((size_t) d_inner_tile * n_t * sizeof(float), 128);
     }
 
     L->src0_bytes = L->src0_bytes_per_thread * n_threads;

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

@@ -0,0 +1,958 @@
+#pragma clang diagnostic ignored "-Wunused-variable"
+
+#include <float.h>
+#include <math.h>
+#include <HAP_farf.h>
+
+#include "hex-common.h"
+#include "dma-queue.h"
+#include "hex-profile.h"
+#include "htp-ctx.h"
+#include "htp-ops.h"
+#include "htp-tensor.h"
+#include "hvx-types.h"
+#include "hvx-base.h"
+#include "hvx-utils.h"
+#include "upscale-ops.h"
+
+static inline int32_t clamp_i32(int32_t v, int32_t lo, int32_t hi) {
+    return v < lo ? lo : (v > hi ? hi : v);
+}
+
+static inline float clamp_f32(float v, float lo, float hi) {
+    return v < lo ? lo : (v > hi ? hi : v);
+}
+
+static inline float bicubic_weight1(float x) {
+    const float a = -0.75f;
+    return ((a + 2.0f) * x - (a + 3.0f)) * x * x + 1.0f;
+}
+
+static inline float bicubic_weight2(float x) {
+    const float a = -0.75f;
+    return ((a * x - 5.0f * a) * x + 8.0f * a) * x - 4.0f * a;
+}
+
+static inline void bicubic_weights(float x, float w[4]) {
+    w[0] = bicubic_weight2(x + 1.0f);
+    w[1] = bicubic_weight1(x);
+    w[2] = bicubic_weight1(1.0f - x);
+    w[3] = bicubic_weight2(2.0f - x);
+}
+
+static void upscale_init_luts_nearest(
+    uint32_t * lut_x,
+    uint32_t * lut_y,
+    const struct htp_tensor * src0,
+    const struct htp_tensor * dst,
+    const struct htp_upscale_kernel_params * p
+) {
+    const uint32_t src_x = src0->ne[0];
+    const uint32_t src_y = src0->ne[1];
+    const uint32_t dst_x = dst->ne[0];
+    const uint32_t dst_y = dst->ne[1];
+    const uint32_t dst_x_aligned = (dst_x + VLEN_FP32 - 1) & ~(VLEN_FP32 - 1);
+
+    for (uint32_t ox = 0; ox < dst_x_aligned; ++ox) {
+        if (ox < dst_x) {
+            uint32_t x0 = (uint32_t) ((float) ox * p->inv_sf0);
+            lut_x[ox] = MIN(x0, src_x - 1) * sizeof(float);
+        } else {
+            lut_x[ox] = 0;
+        }
+    }
+    for (uint32_t oy = 0; oy < dst_y; ++oy) {
+        uint32_t y0 = (uint32_t) ((float) oy * p->inv_sf1);
+        lut_y[oy] = MIN(y0, src_y - 1);
+    }
+}
+
+static void upscale_init_luts_bilinear(
+    struct upscale_bilinear_x_chunk * lut_x,
+    struct upscale_bilinear_y * lut_y,
+    const struct htp_tensor * src0,
+    const struct htp_tensor * dst,
+    const struct htp_upscale_kernel_params * p
+) {
+    const uint32_t src_x = src0->ne[0];
+    const uint32_t src_y = src0->ne[1];
+    const uint32_t dst_x = dst->ne[0];
+    const uint32_t dst_y = dst->ne[1];
+    const uint32_t n_chunks = (dst_x + VLEN_FP32 - 1) / VLEN_FP32;
+
+    for (uint32_t c = 0; c < n_chunks; ++c) {
+        struct upscale_bilinear_x_chunk * chunk = &lut_x[c];
+        const uint32_t ox_start = c * VLEN_FP32;
+        for (uint32_t l = 0; l < VLEN_FP32; ++l) {
+            const uint32_t ox = ox_start + l;
+            if (ox < dst_x) {
+                float xf = ((float) ox + p->pixel_offset) * p->inv_sf0 - p->pixel_offset;
+                int32_t x0 = (int32_t) floorf(xf);
+                int32_t x1 = x0 + 1;
+                x0 = clamp_i32(x0, 0, (int32_t) src_x - 1);
+                x1 = clamp_i32(x1, 0, (int32_t) src_x - 1);
+                float dx = clamp_f32(xf - (float) x0, 0.0f, 1.0f);
+
+                chunk->off0[l] = x0 * (int32_t) sizeof(float);
+                chunk->off1[l] = x1 * (int32_t) sizeof(float);
+                chunk->w0[l]   = 1.0f - dx;
+                chunk->dx[l]   = dx;
+            } else {
+                chunk->off0[l] = 0;
+                chunk->off1[l] = 0;
+                chunk->w0[l]   = 0.0f;
+                chunk->dx[l]   = 0.0f;
+            }
+        }
+    }
+    for (uint32_t oy = 0; oy < dst_y; ++oy) {
+        float yf = ((float) oy + p->pixel_offset) * p->inv_sf1 - p->pixel_offset;
+        int32_t y0 = (int32_t) floorf(yf);
+        int32_t y1 = y0 + 1;
+        y0 = clamp_i32(y0, 0, (int32_t) src_y - 1);
+        y1 = clamp_i32(y1, 0, (int32_t) src_y - 1);
+        float dy = clamp_f32(yf - (float) y0, 0.0f, 1.0f);
+
+        lut_y[oy].y0  = (uint32_t) y0;
+        lut_y[oy].y1  = (uint32_t) y1;
+        lut_y[oy].wy1 = dy;
+        lut_y[oy].wy0 = 1.0f - dy;
+    }
+}
+
+static void upscale_init_luts_bicubic(
+    struct upscale_bicubic_x * lut_x,
+    struct upscale_bicubic_y * lut_y,
+    const struct htp_tensor * src0,
+    const struct htp_tensor * dst,
+    const struct htp_upscale_kernel_params * p
+) {
+    const uint32_t src_x = src0->ne[0];
+    const uint32_t src_y = src0->ne[1];
+    const uint32_t dst_x = dst->ne[0];
+    const uint32_t dst_y = dst->ne[1];
+
+    for (uint32_t ox = 0; ox < dst_x; ++ox) {
+        float xf = ((float) ox + p->pixel_offset) * p->inv_sf0 - p->pixel_offset;
+        int32_t x0 = (int32_t) floorf(xf);
+        float dx = xf - (float) x0;
+        bicubic_weights(dx, lut_x[ox].w);
+
+        lut_x[ox].x[0] = clamp_i32(x0 - 1, 0, (int32_t) src_x - 1);
+        lut_x[ox].x[1] = clamp_i32(x0 + 0, 0, (int32_t) src_x - 1);
+        lut_x[ox].x[2] = clamp_i32(x0 + 1, 0, (int32_t) src_x - 1);
+        lut_x[ox].x[3] = clamp_i32(x0 + 2, 0, (int32_t) src_x - 1);
+    }
+    for (uint32_t oy = 0; oy < dst_y; ++oy) {
+        float yf = ((float) oy + p->pixel_offset) * p->inv_sf1 - p->pixel_offset;
+        int32_t y0 = (int32_t) floorf(yf);
+        float dy = yf - (float) y0;
+        bicubic_weights(dy, lut_y[oy].w);
+
+        lut_y[oy].y[0] = clamp_i32(y0 - 1, 0, (int32_t) src_y - 1);
+        lut_y[oy].y[1] = clamp_i32(y0 + 0, 0, (int32_t) src_y - 1);
+        lut_y[oy].y[2] = clamp_i32(y0 + 1, 0, (int32_t) src_y - 1);
+        lut_y[oy].y[3] = clamp_i32(y0 + 2, 0, (int32_t) src_y - 1);
+    }
+}
+
+static void upscale_init_luts_antialias(
+    struct upscale_antialias_x * lut_x,
+    struct upscale_antialias_y * lut_y,
+    const struct htp_tensor * src0,
+    const struct htp_tensor * dst,
+    const struct htp_upscale_kernel_params * p
+) {
+    const uint32_t src_x = src0->ne[0];
+    const uint32_t src_y = src0->ne[1];
+    const uint32_t dst_x = dst->ne[0];
+    const uint32_t dst_y = dst->ne[1];
+
+    const float filter_radius0 = p->filter_radius0;
+    const float filter_radius1 = p->filter_radius1;
+
+    for (uint32_t ox = 0; ox < dst_x; ++ox) {
+        const float xf = ((float) ox + p->pixel_offset) * p->inv_sf0;
+        int32_t x_min = (int32_t) MAX(xf - filter_radius0 + p->pixel_offset, 0.0f);
+        int32_t x_max = (int32_t) MIN(xf + filter_radius0 + p->pixel_offset, (float) src_x);
+        lut_x[ox].x_min = clamp_i32(x_min, 0, (int32_t) src_x);
+        lut_x[ox].x_max = clamp_i32(x_max, 0, (int32_t) src_x);
+        lut_x[ox].xf    = xf;
+    }
+    for (uint32_t oy = 0; oy < dst_y; ++oy) {
+        const float yf = ((float) oy + p->pixel_offset) * p->inv_sf1;
+        int32_t y_min = (int32_t) MAX(yf - filter_radius1 + p->pixel_offset, 0.0f);
+        int32_t y_max = (int32_t) MIN(yf + filter_radius1 + p->pixel_offset, (float) src_y);
+        lut_y[oy].y_min = clamp_i32(y_min, 0, (int32_t) src_y);
+        lut_y[oy].y_max = clamp_i32(y_max, 0, (int32_t) src_y);
+        lut_y[oy].yf    = yf;
+    }
+}
+
+static inline void get_src_y_window(
+    const struct htp_tensor * src0,
+    const struct htp_tensor * dst,
+    const struct htp_upscale_kernel_params * p,
+    const void * lut_y,
+    uint32_t dst_y_start,
+    uint32_t dst_y_count,
+    uint32_t * out_src_y_start,
+    uint32_t * out_src_y_count
+) {
+    const uint32_t src_y = src0->ne[1];
+    const uint32_t dst_y = dst->ne[1];
+
+    if (dst_y_count == dst_y) {
+        *out_src_y_start = 0;
+        *out_src_y_count = src_y;
+        return;
+    }
+
+    const uint32_t last_oy = dst_y_start + dst_y_count - 1;
+
+    if (p->mode == HTP_UPSCALE_MODE_NEAREST) {
+        const uint32_t * ly = (const uint32_t *) lut_y;
+        const uint32_t y0 = ly[dst_y_start];
+        const uint32_t y1 = ly[last_oy];
+        *out_src_y_start = y0;
+        *out_src_y_count = y1 - y0 + 1;
+    } else if (p->mode == HTP_UPSCALE_MODE_BICUBIC) {
+        const struct upscale_bicubic_y * ly = (const struct upscale_bicubic_y *) lut_y;
+        const int32_t y0 = ly[dst_y_start].y[0];
+        const int32_t y1 = ly[last_oy].y[3];
+        *out_src_y_start = (uint32_t) y0;
+        *out_src_y_count = (uint32_t) (y1 - y0 + 1);
+    } else if (p->antialias) {
+        const struct upscale_antialias_y * ly = (const struct upscale_antialias_y *) lut_y;
+        const int32_t y0 = ly[dst_y_start].y_min;
+        const int32_t y1 = ly[last_oy].y_max;
+        *out_src_y_start = (uint32_t) y0;
+        *out_src_y_count = (uint32_t) (y1 - y0);
+    } else {
+        const struct upscale_bilinear_y * ly = (const struct upscale_bilinear_y *) lut_y;
+        const uint32_t y0 = ly[dst_y_start].y0;
+        const uint32_t y1 = ly[last_oy].y1;
+        *out_src_y_start = y0;
+        *out_src_y_count = y1 - y0 + 1;
+    }
+}
+
+static inline void resolve_dst_slice(
+    const struct htp_tensor * dst,
+    const struct htp_upscale_kernel_params * p,
+    uint32_t chunk_idx,
+    dma_addr_t * out_slice_dst_off,
+    uint32_t * out_dst_ys,
+    uint32_t * out_dst_yc
+) {
+    const uint32_t s_idx = fastdiv(chunk_idx, &p->div_chunks_per_slice);
+    const uint32_t c_idx = chunk_idx - s_idx * p->chunks_per_slice;
+    const uint32_t d_ys  = c_idx * p->dst_nrows_per_chunk;
+    const uint32_t d_yc  = MIN(p->dst_nrows_per_chunk, dst->ne[1] - d_ys);
+    const uint32_t i3    = fastdiv(s_idx, &p->div_dst_z);
+    const uint32_t i2    = s_idx - i3 * dst->ne[2];
+
+    *out_slice_dst_off = (dma_addr_t) i2 * dst->nb[2] + (dma_addr_t) i3 * dst->nb[3];
+    *out_dst_ys        = d_ys;
+    *out_dst_yc        = d_yc;
+}
+
+static inline void resolve_src_slice(
+    const struct htp_tensor * src0,
+    const struct htp_tensor * dst,
+    const struct htp_upscale_kernel_params * p,
+    const void * lut_y,
+    uint32_t chunk_idx,
+    dma_addr_t * out_slice_src_off,
+    uint32_t * out_src_ys,
+    uint32_t * out_src_yc
+) {
+    const uint32_t s_idx = fastdiv(chunk_idx, &p->div_chunks_per_slice);
+    const uint32_t c_idx = chunk_idx - s_idx * p->chunks_per_slice;
+    const uint32_t d_ys  = c_idx * p->dst_nrows_per_chunk;
+    const uint32_t d_yc  = MIN(p->dst_nrows_per_chunk, dst->ne[1] - d_ys);
+    uint32_t s_ys, s_yc;
+    get_src_y_window(src0, dst, p, lut_y, d_ys, d_yc, &s_ys, &s_yc);
+
+    const uint32_t i3  = fastdiv(s_idx, &p->div_dst_z);
+    const uint32_t i2  = s_idx - i3 * dst->ne[2];
+    const uint32_t i03 = (uint32_t) ((float) i3 * p->inv_sf3);
+    const uint32_t i02 = (uint32_t) ((float) i2 * p->inv_sf2);
+
+    *out_slice_src_off = (dma_addr_t) i02 * src0->nb[2] + (dma_addr_t) i03 * src0->nb[3];
+    *out_src_ys        = s_ys;
+    *out_src_yc        = s_yc;
+}
+
+struct upscale_context;
+
+typedef void (*upscale_chunk_fn_t)(
+    const struct upscale_context * ctx,
+    const float * srcb,
+    float * dstb,
+    float * R_buf,
+    float * gather_tmp,
+    uint32_t dst_y_start,
+    uint32_t dst_y_count,
+    uint32_t src_y_start
+);
+
+struct upscale_context {
+    struct htp_ops_context * octx;
+    const struct htp_upscale_kernel_params * kparams;
+    struct htp_upscale_vtcm_layout layout;
+    upscale_chunk_fn_t compute_chunk;
+    const void * lut_x;
+    const void * lut_y;
+    uint32_t chunk_start;
+    uint32_t chunk_count;
+    uint32_t chunks_per_thread;
+};
+
+static void compute_nearest_chunk(
+    const struct upscale_context * ctx,
+    const float * srcb,
+    float * dstb,
+    float * R_buf,
+    float * gather_tmp,
+    uint32_t dst_y_start,
+    uint32_t dst_y_count,
+    uint32_t src_y_start
+) {
+    const uint32_t src_x = ctx->octx->src[0]->ne[0];
+    const uint32_t dst_x = ctx->octx->dst->ne[0];
+    const uint32_t * lut_x = (const uint32_t *) ctx->lut_x;
+    const uint32_t * lut_y = (const uint32_t *) ctx->lut_y;
+    const uint32_t region_size = hex_round_up(src_x * sizeof(float), 128);
+
+    for (uint32_t oy = dst_y_start; oy < dst_y_start + dst_y_count; ++oy) {
+        const uint32_t y0 = lut_y[oy];
+        const float * r0  = srcb + (y0 - src_y_start) * src_x;
+        float * dst_row   = dstb + (oy - dst_y_start) * dst_x;
+
+        for (uint32_t ix = 0; ix < src_x; ix += VLEN_FP32) {
+            *(HVX_Vector *) (R_buf + ix) = hvx_vmemu(r0 + ix);
+        }
+
+        for (uint32_t ox = 0; ox < dst_x; ox += VLEN_FP32) {
+            const uint32_t rem = dst_x - ox;
+            const uint32_t nbytes = (rem < VLEN_FP32) ? (rem * sizeof(float)) : VLEN;
+            const HVX_Vector voff = hvx_vmem(lut_x + ox);
+            Q6_vgather_ARMVw((HVX_Vector *) gather_tmp, (size_t) R_buf, region_size, voff);
+            hvx_gather_sync(gather_tmp);
+            hvx_vec_store_u(dst_row + ox, nbytes, hvx_vmem(gather_tmp));
+        }
+    }
+}
+
+static void compute_nearest_2x_chunk(
+    const struct upscale_context * ctx,
+    const float * srcb,
+    float * dstb,
+    float * R_buf,
+    float * gather_tmp,
+    uint32_t dst_y_start,
+    uint32_t dst_y_count,
+    uint32_t src_y_start
+) {
+    (void) R_buf;
+    (void) gather_tmp;
+    const uint32_t src_x = ctx->octx->src[0]->ne[0];
+    const uint32_t dst_x = ctx->octx->dst->ne[0];
+
+    for (uint32_t oy = dst_y_start; oy < dst_y_start + dst_y_count; ++oy) {
+        const uint32_t sy = oy >> 1;
+        const float * r0  = srcb + (sy - src_y_start) * src_x;
+        float * dst_row   = dstb + (oy - dst_y_start) * dst_x;
+
+        for (uint32_t ix = 0; ix < src_x; ix += VLEN_FP32) {
+            const HVX_Vector vin = hvx_vmemu(r0 + ix);
+
+            const HVX_VectorPair p = Q6_W_vshuff_VVR(vin, vin, -4);
+            const HVX_Vector v_lo = Q6_V_lo_W(p);
+            const HVX_Vector v_hi = Q6_V_hi_W(p);
+
+            const uint32_t ox0 = 2 * ix;
+            const uint32_t rem_out0 = (dst_x > ox0) ? (dst_x - ox0) : 0;
+            const uint32_t nbytes_out0 = (rem_out0 < VLEN_FP32) ? (rem_out0 * sizeof(float)) : VLEN;
+            hvx_vec_store_u(dst_row + ox0, nbytes_out0, v_lo);
+
+            const uint32_t ox1 = ox0 + VLEN_FP32;
+            if (dst_x > ox1) {
+                const uint32_t rem_out1 = dst_x - ox1;
+                const uint32_t nbytes_out1 = (rem_out1 < VLEN_FP32) ? (rem_out1 * sizeof(float)) : VLEN;
+                hvx_vec_store_u(dst_row + ox1, nbytes_out1, v_hi);
+            }
+        }
+    }
+}
+
+static void compute_nearest_4x_chunk(
+    const struct upscale_context * ctx,
+    const float * srcb,
+    float * dstb,
+    float * R_buf,
+    float * gather_tmp,
+    uint32_t dst_y_start,
+    uint32_t dst_y_count,
+    uint32_t src_y_start
+) {
+    (void) R_buf;
+    (void) gather_tmp;
+    const uint32_t src_x = ctx->octx->src[0]->ne[0];
+    const uint32_t dst_x = ctx->octx->dst->ne[0];
+
+    for (uint32_t oy = dst_y_start; oy < dst_y_start + dst_y_count; ++oy) {
+        const uint32_t sy = oy >> 2;
+        const float * r0  = srcb + (sy - src_y_start) * src_x;
+        float * dst_row   = dstb + (oy - dst_y_start) * dst_x;
+
+        for (uint32_t ix = 0; ix < src_x; ix += VLEN_FP32) {
+            const HVX_Vector vin = hvx_vmemu(r0 + ix);
+
+            const HVX_VectorPair p = Q6_W_vshuff_VVR(vin, vin, -4);
+            const HVX_Vector v_lo = Q6_V_lo_W(p);
+            const HVX_Vector v_hi = Q6_V_hi_W(p);
+
+            const HVX_VectorPair p0 = Q6_W_vshuff_VVR(v_lo, v_lo, -4);
+            const HVX_VectorPair p1 = Q6_W_vshuff_VVR(v_hi, v_hi, -4);
+
+            const HVX_Vector out0 = Q6_V_lo_W(p0);
+            const HVX_Vector out1 = Q6_V_hi_W(p0);
+            const HVX_Vector out2 = Q6_V_lo_W(p1);
+            const HVX_Vector out3 = Q6_V_hi_W(p1);
+
+            const uint32_t ox = 4 * ix;
+            if (dst_x > ox) {
+                const uint32_t rem = dst_x - ox;
+                hvx_vec_store_u(dst_row + ox, (rem < VLEN_FP32) ? (rem * sizeof(float)) : VLEN, out0);
+            }
+            if (dst_x > ox + VLEN_FP32) {
+                const uint32_t rem = dst_x - (ox + VLEN_FP32);
+                hvx_vec_store_u(dst_row + ox + VLEN_FP32, (rem < VLEN_FP32) ? (rem * sizeof(float)) : VLEN, out1);
+            }
+            if (dst_x > ox + 2 * VLEN_FP32) {
+                const uint32_t rem = dst_x - (ox + 2 * VLEN_FP32);
+                hvx_vec_store_u(dst_row + ox + 2 * VLEN_FP32, (rem < VLEN_FP32) ? (rem * sizeof(float)) : VLEN, out2);
+            }
+            if (dst_x > ox + 3 * VLEN_FP32) {
+                const uint32_t rem = dst_x - (ox + 3 * VLEN_FP32);
+                hvx_vec_store_u(dst_row + ox + 3 * VLEN_FP32, (rem < VLEN_FP32) ? (rem * sizeof(float)) : VLEN, out3);
+            }
+        }
+    }
+}
+
+static inline void compute_bilinear_chunk_32(
+    float * restrict dst_chunk,
+    const float * restrict R_buf,
+    float * restrict tmp_a,
+    float * restrict tmp_b,
+    uint32_t region_size,
+    const struct upscale_bilinear_x_chunk * restrict chunk,
+    uint32_t nbytes
+) {
+    const HVX_Vector voff0 = hvx_vmem(chunk->off0);
+    const HVX_Vector voff1 = hvx_vmem(chunk->off1);
+    const HVX_Vector vw0   = hvx_vmem(chunk->w0);
+    const HVX_Vector vdx   = hvx_vmem(chunk->dx);
+
+    Q6_vgather_ARMVw((HVX_Vector *) tmp_a, (size_t) R_buf, region_size, voff0);
+    Q6_vgather_ARMVw((HVX_Vector *) tmp_b, (size_t) R_buf, region_size, voff1);
+    hvx_gather_sync(tmp_b);
+
+    const HVX_Vector va = hvx_vmem(tmp_a);
+    const HVX_Vector vb = hvx_vmem(tmp_b);
+
+    const HVX_Vector val = hvx_vec_add_f32_f32(
+        hvx_vec_mul_f32_f32(vw0, va),
+        hvx_vec_mul_f32_f32(vdx, vb)
+    );
+
+    hvx_vec_store_u(dst_chunk, nbytes, val);
+}
+
+static void compute_bilinear_chunk(
+    const struct upscale_context * ctx,
+    const float * srcb,
+    float * dstb,
+    float * R_buf,
+    float * gather_tmp,
+    uint32_t dst_y_start,
+    uint32_t dst_y_count,
+    uint32_t src_y_start
+) {
+    const uint32_t src_x = ctx->octx->src[0]->ne[0];
+    const uint32_t dst_x = ctx->octx->dst->ne[0];
+    const struct upscale_bilinear_x_chunk * lut_x = (const struct upscale_bilinear_x_chunk *) ctx->lut_x;
+    const struct upscale_bilinear_y * lut_y = (const struct upscale_bilinear_y *) ctx->lut_y;
+    const uint32_t n_chunks = (dst_x + VLEN_FP32 - 1) / VLEN_FP32;
+    const uint32_t region_size = hex_round_up(src_x * sizeof(float), 128);
+    float * tmp_a = gather_tmp;
+    float * tmp_b = gather_tmp + VLEN_FP32;
+
+    for (uint32_t oy = dst_y_start; oy < dst_y_start + dst_y_count; ++oy) {
+        const struct upscale_bilinear_y * ey = &lut_y[oy];
+        const float * r0 = srcb + (ey->y0 - src_y_start) * src_x;
+        const float * r1 = srcb + (ey->y1 - src_y_start) * src_x;
+        float * dst_row  = dstb + (oy - dst_y_start) * dst_x;
+
+        if (ey->y0 == ey->y1) {
+            for (uint32_t ix = 0; ix < src_x; ix += VLEN_FP32) {
+                *(HVX_Vector *) (R_buf + ix) = hvx_vmemu(r0 + ix);
+            }
+        } else {
+            const HVX_Vector vwy0 = hvx_vec_splat_f32(ey->wy0);
+            const HVX_Vector vwy1 = hvx_vec_splat_f32(ey->wy1);
+            for (uint32_t ix = 0; ix < src_x; ix += VLEN_FP32) {
+                const HVX_Vector vr0 = hvx_vmemu(r0 + ix);
+                const HVX_Vector vr1 = hvx_vmemu(r1 + ix);
+                *(HVX_Vector *) (R_buf + ix) = hvx_vec_add_f32_f32(
+                    hvx_vec_mul_f32_f32(vwy0, vr0),
+                    hvx_vec_mul_f32_f32(vwy1, vr1)
+                );
+            }
+        }
+
+        for (uint32_t c = 0; c < n_chunks; ++c) {
+            const uint32_t ox = c * VLEN_FP32;
+            const uint32_t rem = dst_x - ox;
+            const uint32_t nbytes = (rem < VLEN_FP32) ? (rem * sizeof(float)) : VLEN;
+            compute_bilinear_chunk_32(dst_row + ox, R_buf, tmp_a, tmp_b, region_size, &lut_x[c], nbytes);
+        }
+    }
+}
+
+static void compute_bilinear_2x_chunk(
+    const struct upscale_context * ctx,
+    const float * srcb,
+    float * dstb,
+    float * R_buf,
+    float * gather_tmp,
+    uint32_t dst_y_start,
+    uint32_t dst_y_count,
+    uint32_t src_y_start
+) {
+    (void) R_buf;
+    (void) gather_tmp;
+    const uint32_t src_x = ctx->octx->src[0]->ne[0];
+    const uint32_t dst_x = ctx->octx->dst->ne[0];
+    const struct upscale_bilinear_y * lut_y = (const struct upscale_bilinear_y *) ctx->lut_y;
+    const HVX_Vector c_075 = hvx_vec_splat_f32(0.75f);
+    const HVX_Vector c_025 = hvx_vec_splat_f32(0.25f);
+
+    for (uint32_t oy = dst_y_start; oy < dst_y_start + dst_y_count; ++oy) {
+        const struct upscale_bilinear_y * ey = &lut_y[oy];
+        const float * r0 = srcb + (ey->y0 - src_y_start) * src_x;
+        const float * r1 = srcb + (ey->y1 - src_y_start) * src_x;
+        float * dst_row  = dstb + (oy - dst_y_start) * dst_x;
+
+        const bool same_row = (ey->y0 == ey->y1);
+        const HVX_Vector vwy0 = hvx_vec_splat_f32(ey->wy0);
+        const HVX_Vector vwy1 = hvx_vec_splat_f32(ey->wy1);
+
+        HVX_Vector v_curr = hvx_vmem(r0);
+        if (!same_row) {
+            const HVX_Vector v1 = hvx_vmem(r1);
+            v_curr = hvx_vec_add_f32_f32(hvx_vec_mul_f32_f32(vwy0, v_curr), hvx_vec_mul_f32_f32(vwy1, v1));
+        }
+
+        HVX_Vector v_prev = Q6_V_vror_VR(v_curr, 4);
+
+        for (uint32_t ix = 0; ix < src_x; ix += VLEN_FP32) {
+            HVX_Vector v_next_chunk;
+            const uint32_t next_ix = ix + VLEN_FP32;
+            if (next_ix < src_x) {
+                v_next_chunk = hvx_vmem(r0 + next_ix);
+                if (!same_row) {
+                    const HVX_Vector v1 = hvx_vmem(r1 + next_ix);
+                    v_next_chunk = hvx_vec_add_f32_f32(hvx_vec_mul_f32_f32(vwy0, v_next_chunk), hvx_vec_mul_f32_f32(vwy1, v1));
+                }
+            } else {
+                v_next_chunk = Q6_V_vror_VR(v_curr, 124);
+            }
+
+            const HVX_Vector v_shl = Q6_V_valign_VVR(v_curr, v_prev, 124);
+            const HVX_Vector v_shr = Q6_V_valign_VVR(v_next_chunk, v_curr, 4);
+
+            const HVX_Vector ve = hvx_vec_add_f32_f32(hvx_vec_mul_f32_f32(c_075, v_curr), hvx_vec_mul_f32_f32(c_025, v_shl));
+            const HVX_Vector vo = hvx_vec_add_f32_f32(hvx_vec_mul_f32_f32(c_075, v_curr), hvx_vec_mul_f32_f32(c_025, v_shr));
+
+            const HVX_VectorPair p = Q6_W_vshuff_VVR(vo, ve, -4);
+            const HVX_Vector out0 = Q6_V_lo_W(p);
+            const HVX_Vector out1 = Q6_V_hi_W(p);
+
+            const uint32_t ox0 = 2 * ix;
+            hvx_vmem(dst_row + ox0)             = out0;
+            hvx_vmem(dst_row + ox0 + VLEN_FP32) = out1;
+
+            v_prev = v_curr;
+            v_curr = v_next_chunk;
+        }
+    }
+}
+
+static void compute_bilinear_antialias_chunk(
+    const struct upscale_context * ctx,
+    const float * srcb,
+    float * dstb,
+    float * R_buf,
+    float * gather_tmp,
+    uint32_t dst_y_start,
+    uint32_t dst_y_count,
+    uint32_t src_y_start
+) {
+    (void) R_buf;
+    (void) gather_tmp;
+    const uint32_t src_x = ctx->octx->src[0]->ne[0];
+    const uint32_t dst_x = ctx->octx->dst->ne[0];
+    const struct htp_upscale_kernel_params * p = ctx->kparams;
+    const struct upscale_antialias_x * lut_x = (const struct upscale_antialias_x *) ctx->lut_x;
+    const struct upscale_antialias_y * lut_y = (const struct upscale_antialias_y *) ctx->lut_y;
+
+    const float invscale1 = p->invscale1;
+    const float invscale0 = p->invscale0;
+
+    for (uint32_t oy = dst_y_start; oy < dst_y_start + dst_y_count; ++oy) {
+        const struct upscale_antialias_y * ey = &lut_y[oy];
+        const float yf = ey->yf;
+        const int32_t y_min = ey->y_min;
+        const int32_t y_max = ey->y_max;
+
+        float * dst_row = dstb + (oy - dst_y_start) * dst_x;
+
+        for (uint32_t ox = 0; ox < dst_x; ox += VLEN_FP32) {
+            const uint32_t rem = dst_x - ox;
+            const uint32_t nbytes = (rem < VLEN_FP32) ? (rem * sizeof(float)) : VLEN;
+            float tmp[VLEN_FP32] __attribute__((aligned(128)));
+            const uint32_t cnt = rem < VLEN_FP32 ? rem : VLEN_FP32;
+
+            for (uint32_t l = 0; l < cnt; ++l) {
+                const uint32_t cur_ox = ox + l;
+                const struct upscale_antialias_x * ex = &lut_x[cur_ox];
+                const float xf = ex->xf;
+                const int32_t x_min = ex->x_min;
+                const int32_t x_max = ex->x_max;
+
+                float val = 0.0f;
+                float total_weight = 0.0f;
+
+                for (int32_t sy = y_min; sy < y_max; ++sy) {
+                    const float diff_y = fabsf((float) sy - yf + p->pixel_offset) * invscale1;
+                    const float weight_y = MAX(1.0f - diff_y, 0.0f);
+                    if (weight_y <= 0.0f) {
+                        continue;
+                    }
+                    const float * r = srcb + (uint32_t) (sy - (int32_t) src_y_start) * src_x;
+
+                    for (int32_t sx = x_min; sx < x_max; ++sx) {
+                        const float diff_x = fabsf((float) sx - xf + p->pixel_offset) * invscale0;
+                        const float weight_x = MAX(1.0f - diff_x, 0.0f);
+                        const float weight = weight_x * weight_y;
+                        if (weight <= 0.0f) {
+                            continue;
+                        }
+                        val += r[sx] * weight;
+                        total_weight += weight;
+                    }
+                }
+
+                tmp[l] = (total_weight > 0.0f) ? (val / total_weight) : 0.0f;
+            }
+
+            hvx_vec_store_u(dst_row + ox, nbytes, *(const HVX_Vector *) tmp);
+        }
+    }
+}
+
+static void compute_bicubic_chunk(
+    const struct upscale_context * ctx,
+    const float * srcb,
+    float * dstb,
+    float * R_buf,
+    float * gather_tmp,
+    uint32_t dst_y_start,
+    uint32_t dst_y_count,
+    uint32_t src_y_start
+) {
+    (void) R_buf;
+    (void) gather_tmp;
+    const uint32_t src_x = ctx->octx->src[0]->ne[0];
+    const uint32_t dst_x = ctx->octx->dst->ne[0];
+    const struct upscale_bicubic_x * lut_x = (const struct upscale_bicubic_x *) ctx->lut_x;
+    const struct upscale_bicubic_y * lut_y = (const struct upscale_bicubic_y *) ctx->lut_y;
+
+    for (uint32_t oy = dst_y_start; oy < dst_y_start + dst_y_count; ++oy) {
+        const struct upscale_bicubic_y * ey = &lut_y[oy];
+        const float * r0 = srcb + (uint32_t) (ey->y[0] - (int32_t) src_y_start) * src_x;
+        const float * r1 = srcb + (uint32_t) (ey->y[1] - (int32_t) src_y_start) * src_x;
+        const float * r2 = srcb + (uint32_t) (ey->y[2] - (int32_t) src_y_start) * src_x;
+        const float * r3 = srcb + (uint32_t) (ey->y[3] - (int32_t) src_y_start) * src_x;
+
+        float * dst_row = dstb + (oy - dst_y_start) * dst_x;
+
+        const HVX_Vector vwy0 = hvx_vec_splat_f32(ey->w[0]);
+        const HVX_Vector vwy1 = hvx_vec_splat_f32(ey->w[1]);
+        const HVX_Vector vwy2 = hvx_vec_splat_f32(ey->w[2]);
+        const HVX_Vector vwy3 = hvx_vec_splat_f32(ey->w[3]);
+
+        for (uint32_t ox = 0; ox < dst_x; ox += VLEN_FP32) {
+            const uint32_t rem = dst_x - ox;
+            const uint32_t nbytes = (rem < VLEN_FP32) ? (rem * sizeof(float)) : VLEN;
+
+            float v0[VLEN_FP32] __attribute__((aligned(128)));
+            float v1[VLEN_FP32] __attribute__((aligned(128)));
+            float v2[VLEN_FP32] __attribute__((aligned(128)));
+            float v3[VLEN_FP32] __attribute__((aligned(128)));
+
+            const uint32_t cnt = rem < VLEN_FP32 ? rem : VLEN_FP32;
+            for (uint32_t l = 0; l < cnt; ++l) {
+                const struct upscale_bicubic_x * e = &lut_x[ox + l];
+                v0[l] = r0[e->x[0]] * e->w[0] + r0[e->x[1]] * e->w[1] + r0[e->x[2]] * e->w[2] + r0[e->x[3]] * e->w[3];
+                v1[l] = r1[e->x[0]] * e->w[0] + r1[e->x[1]] * e->w[1] + r1[e->x[2]] * e->w[2] + r1[e->x[3]] * e->w[3];
+                v2[l] = r2[e->x[0]] * e->w[0] + r2[e->x[1]] * e->w[1] + r2[e->x[2]] * e->w[2] + r2[e->x[3]] * e->w[3];
+                v3[l] = r3[e->x[0]] * e->w[0] + r3[e->x[1]] * e->w[1] + r3[e->x[2]] * e->w[2] + r3[e->x[3]] * e->w[3];
+            }
+            for (uint32_t l = cnt; l < VLEN_FP32; ++l) {
+                v0[l] = v1[l] = v2[l] = v3[l] = 0.0f;
+            }
+
+            const HVX_Vector vv0 = *(const HVX_Vector *) v0;
+            const HVX_Vector vv1 = *(const HVX_Vector *) v1;
+            const HVX_Vector vv2 = *(const HVX_Vector *) v2;
+            const HVX_Vector vv3 = *(const HVX_Vector *) v3;
+
+            const HVX_Vector acc0 = hvx_vec_add_f32_f32(hvx_vec_mul_f32_f32(vwy0, vv0), hvx_vec_mul_f32_f32(vwy1, vv1));
+            const HVX_Vector acc1 = hvx_vec_add_f32_f32(hvx_vec_mul_f32_f32(vwy2, vv2), hvx_vec_mul_f32_f32(vwy3, vv3));
+            const HVX_Vector val  = hvx_vec_add_f32_f32(acc0, acc1);
+
+            hvx_vec_store_u(dst_row + ox, nbytes, val);
+        }
+    }
+}
+
+static inline void dma_push_src_chunk(
+    dma_queue * dma_q,
+    float * spad,
+    const struct htp_tensor * src0,
+    dma_addr_t slice_src_off,
+    uint32_t src_ys,
+    uint32_t src_yc
+) {
+    const uint32_t src_x = src0->ne[0];
+    const dma_addr_t src_addr = src0->data + slice_src_off + (dma_addr_t) src_ys * src0->nb[1];
+
+    if (src0->nb[0] == sizeof(float)) {
+        const uint32_t row_size   = src_x * sizeof(float);
+        const uint32_t dst_stride = row_size;
+        const uint32_t src_stride = src0->nb[1];
+        dma_queue_push(dma_q, dma_make_data(spad, src_addr), dst_stride, src_stride, row_size, src_yc);
+    } else {
+        const uint32_t row_size   = sizeof(float);
+        const uint32_t dst_stride = sizeof(float);
+        const uint32_t src_stride = src0->nb[0];
+        const uint32_t nrows      = src_x * src_yc;
+        dma_queue_push(dma_q, dma_make_data(spad, src_addr), dst_stride, src_stride, row_size, nrows);
+    }
+}
+
+static void upscale_worker(unsigned int nth, unsigned int ith, void * data) {
+    struct upscale_context * ctx = (struct upscale_context *) data;
+    const struct htp_upscale_kernel_params * p = ctx->kparams;
+    const struct htp_upscale_vtcm_layout * L = &ctx->layout;
+    const struct htp_tensor * src0 = ctx->octx->src[0];
+    const struct htp_tensor * dst  = ctx->octx->dst;
+
+    const uint32_t first = ctx->chunk_start + ith * ctx->chunks_per_thread;
+    const uint32_t last  = MIN(first + ctx->chunks_per_thread, ctx->chunk_start + ctx->chunk_count);
+
+    if (first >= last) {
+        return;
+    }
+
+    struct htp_thread_trace * tr = &ctx->octx->ctx->trace[ith];
+    dma_queue * dma_q            = ctx->octx->ctx->dma[ith];
+
+    const uint32_t src_spad_half = (uint32_t) L->src_spad_half_size;
+    const uint32_t dst_spad_half = (uint32_t) L->dst_spad_half_size;
+
+    uint8_t * vtcm_base  = (uint8_t *) ctx->octx->ctx->vtcm_base;
+    uint8_t * src_spad   = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, L->off_src)    + ith * L->src_bytes_per_thread;
+    uint8_t * dst_spad   = VTCM_LAYOUT_PTR(uint8_t, vtcm_base, L->off_dst)    + ith * L->dst_bytes_per_thread;
+    float   * R_buf      = VTCM_LAYOUT_PTR(float,   vtcm_base, L->off_r)      + ith * (L->r_bytes_per_thread / sizeof(float));
+    float   * gather_tmp = VTCM_LAYOUT_PTR(float,   vtcm_base, L->off_gather) + ith * (L->gather_bytes_per_thread / sizeof(float));
+    const void * lut_y   = ctx->lut_y;
+
+    float * srcb2[2] = { (float *) src_spad, (float *) (src_spad + src_spad_half) };
+    float * dstb2[2] = { (float *) dst_spad, (float *) (dst_spad + dst_spad_half) };
+
+    const uint32_t dst_row_size = dst->ne[0] * sizeof(float);
+    const uint32_t dst_stride   = dst->nb[1];
+
+    const uint32_t total = last - first;
+
+    // Warmup: prime pipeline with initial transfers
+    for (uint32_t i = 0; i < total && i < 2; ++i) {
+        dma_queue_push(dma_q, dma_make_data(dst->data, dstb2[i]), dst_stride, dst_row_size, dst_row_size, 0);
+
+        dma_addr_t slice_src_off;
+        uint32_t src_ys, src_yc;
+        resolve_src_slice(src0, dst, p, lut_y, first + i, &slice_src_off, &src_ys, &src_yc);
+
+        dma_push_src_chunk(dma_q, srcb2[i], src0, slice_src_off, src_ys, src_yc);
+    }
+
+    // Steady state loop
+    for (uint32_t i = 0; i < total; ++i) {
+        const uint32_t buf = i & 1u;
+        float * srcb = srcb2[buf];
+        float * dstb = dstb2[buf];
+
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+
+        dma_addr_t slice_dst_off;
+        uint32_t dst_ys, dst_yc;
+        resolve_dst_slice(dst, p, first + i, &slice_dst_off, &dst_ys, &dst_yc);
+
+        uint32_t src_ys, src_yc;
+        get_src_y_window(src0, dst, p, lut_y, dst_ys, dst_yc, &src_ys, &src_yc);
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) (first + i));
+        ctx->compute_chunk(ctx, srcb, dstb, R_buf, gather_tmp, dst_ys, dst_yc, src_ys);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) (first + i));
+
+        const dma_addr_t dst_addr = dst->data + slice_dst_off + (dma_addr_t) dst_ys * dst_stride;
+        dma_queue_push(dma_q, dma_make_data(dst_addr, dstb), dst_stride, dst_row_size, dst_row_size, dst_yc);
+
+        const uint32_t next = i + 2;
+        if (next < total) {
+            dma_addr_t next_src_off;
+            uint32_t next_src_ys, next_src_yc;
+            resolve_src_slice(src0, dst, p, lut_y, first + next, &next_src_off, &next_src_ys, &next_src_yc);
+
+            dma_push_src_chunk(dma_q, srcb, src0, next_src_off, next_src_ys, next_src_yc);
+        }
+    }
+
+    dma_queue_flush(dma_q);
+
+    (void) nth;
+}
+
+int op_upscale(struct htp_ops_context * octx) {
+    const struct htp_tensor * src0 = octx->src[0];
+    const struct htp_tensor * dst  = octx->dst;
+    const struct htp_upscale_kernel_params * p =
+        (const struct htp_upscale_kernel_params *) octx->kernel_params;
+
+    const uint32_t src_x = src0->ne[0];
+    const uint32_t src_y = src0->ne[1];
+    const uint32_t dst_x = dst->ne[0];
+    const uint32_t dst_y = dst->ne[1];
+
+    uint32_t chunk_start = 0;
+    uint32_t chunk_count = p->total_chunks;
+
+    if (octx->ctx->mdev.count > 1) {
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(
+            chunk_count, htp_tensor_mdev_data_aligned(dst) ? p->chunks_per_group : 0,
+            octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        chunk_start = range.start;
+        chunk_count = range.count;
+    }
+
+    if (chunk_count == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    const uint32_t n_threads = MIN(p->n_threads, chunk_count);
+    if (!htp_ops_context_set_n_threads(octx, n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    const uint32_t chunks_per_thread = fastdiv(chunk_count + n_threads - 1, &octx->n_threads_div);
+
+    size_t lut_x_size = 0;
+    size_t elem_y_sz  = sizeof(uint32_t);
+
+    if (p->mode == HTP_UPSCALE_MODE_BICUBIC) {
+        lut_x_size = hex_round_up(dst_x * sizeof(struct upscale_bicubic_x), 128);
+        elem_y_sz  = sizeof(struct upscale_bicubic_y);
+    } else if (p->antialias) {
+        lut_x_size = hex_round_up(dst_x * sizeof(struct upscale_antialias_x), 128);
+        elem_y_sz  = sizeof(struct upscale_antialias_y);
+    } else if (p->mode == HTP_UPSCALE_MODE_BILINEAR) {
+        const uint32_t n_chunks = (dst_x + VLEN_FP32 - 1) / VLEN_FP32;
+        lut_x_size = hex_round_up(n_chunks * sizeof(struct upscale_bilinear_x_chunk), 128);
+        elem_y_sz  = sizeof(struct upscale_bilinear_y);
+    } else {
+        lut_x_size = hex_round_up(dst_x * sizeof(uint32_t), 128);
+    }
+
+    const size_t lut_y_size = hex_round_up(dst_y * elem_y_sz, 128);
+    const size_t total_lut_size = lut_x_size + lut_y_size;
+
+    if (total_lut_size > octx->ctx->ddr_spad_size) {
+        return HTP_STATUS_NO_SUPPORT;
+    }
+
+    void * base_scratch = octx->ctx->ddr_spad_base;
+    void * lut_x        = base_scratch;
+    void * lut_y        = (uint8_t *) base_scratch + lut_x_size;
+
+    if (p->mode == HTP_UPSCALE_MODE_NEAREST) {
+        upscale_init_luts_nearest((uint32_t *) lut_x, (uint32_t *) lut_y, src0, dst, p);
+    } else if (p->mode == HTP_UPSCALE_MODE_BICUBIC) {
+        upscale_init_luts_bicubic((struct upscale_bicubic_x *) lut_x, (struct upscale_bicubic_y *) lut_y, src0, dst, p);
+    } else if (p->antialias) {
+        upscale_init_luts_antialias((struct upscale_antialias_x *) lut_x, (struct upscale_antialias_y *) lut_y, src0, dst, p);
+    } else {
+        upscale_init_luts_bilinear((struct upscale_bilinear_x_chunk *) lut_x, (struct upscale_bilinear_y *) lut_y, src0, dst, p);
+    }
+
+    upscale_chunk_fn_t compute_chunk;
+    switch (p->kernel_type) {
+        case HTP_UPSCALE_KERNEL_NEAREST:
+            compute_chunk = compute_nearest_chunk;
+            break;
+        case HTP_UPSCALE_KERNEL_NEAREST_2X:
+            compute_chunk = compute_nearest_2x_chunk;
+            break;
+        case HTP_UPSCALE_KERNEL_NEAREST_4X:
+            compute_chunk = compute_nearest_4x_chunk;
+            break;
+        case HTP_UPSCALE_KERNEL_BILINEAR:
+            compute_chunk = compute_bilinear_chunk;
+            break;
+        case HTP_UPSCALE_KERNEL_BILINEAR_2X:
+            compute_chunk = compute_bilinear_2x_chunk;
+            break;
+        case HTP_UPSCALE_KERNEL_BILINEAR_AA:
+            compute_chunk = compute_bilinear_antialias_chunk;
+            break;
+        case HTP_UPSCALE_KERNEL_BICUBIC:
+            compute_chunk = compute_bicubic_chunk;
+            break;
+        default:
+            return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    struct upscale_context ctx = {
+        .octx               = octx,
+        .kparams            = p,
+        .compute_chunk      = compute_chunk,
+        .lut_x              = lut_x,
+        .lut_y              = lut_y,
+        .chunk_start        = chunk_start,
+        .chunk_count        = chunk_count,
+        .chunks_per_thread  = chunks_per_thread,
+    };
+
+    htp_upscale_vtcm_layout_build(&ctx.layout, src_x, src_y, dst_x, dst_y,
+                                  p->dst_nrows_per_chunk, p->inv_sf1, p->filter_radius1, n_threads);
+
+    work_queue_run(octx->ctx->work_queue, upscale_worker, &ctx, n_threads);
+
+    return HTP_STATUS_OK;
+}

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

@@ -0,0 +1,204 @@
+#ifndef HTP_UPSCALE_OPS_H
+#define HTP_UPSCALE_OPS_H
+
+#include <stdint.h>
+#include <stddef.h>
+#include <stdbool.h>
+#include <math.h>
+
+#include "hex-common.h"
+#include "hex-fastdiv.h"
+#include "htp-vtcm.h"
+
+enum htp_upscale_mode {
+    HTP_UPSCALE_MODE_NEAREST  = 0,
+    HTP_UPSCALE_MODE_BILINEAR = 1,
+    HTP_UPSCALE_MODE_BICUBIC  = 2,
+};
+
+enum htp_upscale_kernel_type {
+    HTP_UPSCALE_KERNEL_NEAREST     = 0,
+    HTP_UPSCALE_KERNEL_NEAREST_2X  = 1,
+    HTP_UPSCALE_KERNEL_NEAREST_4X  = 2,
+    HTP_UPSCALE_KERNEL_BILINEAR    = 3,
+    HTP_UPSCALE_KERNEL_BILINEAR_2X = 4,
+    HTP_UPSCALE_KERNEL_BILINEAR_AA = 5,
+    HTP_UPSCALE_KERNEL_BICUBIC     = 6,
+};
+
+struct htp_upscale_kernel_params {
+    uint8_t  kernel_type;
+    uint8_t  mode;
+    uint8_t  align_corners;
+    uint8_t  antialias;
+
+    float    inv_sf0;
+    float    inv_sf1;
+    float    inv_sf2;
+    float    inv_sf3;
+    float    pixel_offset;
+
+    float    filter_radius0;
+    float    filter_radius1;
+    float    invscale0;
+    float    invscale1;
+
+    uint32_t n_threads;
+    uint32_t chunks_per_group;
+
+    uint32_t dst_nrows_per_chunk;
+    uint32_t chunks_per_slice;
+    uint32_t total_chunks;
+    uint32_t vtcm_size;
+
+    struct fastdiv_values div_chunks_per_slice;
+    struct fastdiv_values div_dst_z;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_upscale_kernel_params) <= 128, "htp_upscale_kernel_params is too large");
+#else
+_Static_assert(sizeof(struct htp_upscale_kernel_params) <= 128, "htp_upscale_kernel_params is too large");
+#endif
+
+struct __attribute__((aligned(128))) upscale_bilinear_x_chunk {
+    int32_t  off0[32];
+    int32_t  off1[32];
+    float    w0[32];
+    float    dx[32];
+};
+
+struct upscale_bilinear_y {
+    uint32_t y0;
+    uint32_t y1;
+    float    wy0;
+    float    wy1;
+};
+
+struct upscale_bicubic_x {
+    int32_t x[4];
+    float   w[4];
+};
+
+struct upscale_bicubic_y {
+    int32_t y[4];
+    float   w[4];
+};
+
+struct upscale_antialias_x {
+    int32_t x_min;
+    int32_t x_max;
+    float   xf;
+};
+
+struct upscale_antialias_y {
+    int32_t y_min;
+    int32_t y_max;
+    float   yf;
+};
+
+struct htp_upscale_vtcm_layout {
+    size_t   total_bytes;
+    size_t   off_src;
+    size_t   off_dst;
+    size_t   off_r;
+    size_t   off_gather;
+    size_t   src_bytes_per_thread;
+    size_t   dst_bytes_per_thread;
+    size_t   r_bytes_per_thread;
+    size_t   gather_bytes_per_thread;
+    size_t   src_spad_half_size;
+    size_t   dst_spad_half_size;
+    uint32_t max_src_chunk_rows;
+    uint32_t dst_chunk_rows;
+};
+
+static inline void htp_upscale_vtcm_layout_build(
+    struct htp_upscale_vtcm_layout * L,
+    uint32_t src_x,
+    uint32_t src_y,
+    uint32_t dst_x,
+    uint32_t dst_y,
+    uint32_t dst_chunk_rows,
+    float    inv_sf1,
+    float    filter_radius1,
+    uint32_t n_threads
+) {
+    uint32_t max_src_rows;
+    if (dst_chunk_rows >= dst_y) {
+        max_src_rows = src_y;
+    } else {
+        const float est = (float) dst_chunk_rows * inv_sf1 + 2.0f * filter_radius1;
+        max_src_rows = (uint32_t) ceilf(est) + 4;
+        if (max_src_rows > src_y) {
+            max_src_rows = src_y;
+        }
+    }
+
+    const uint32_t src_chunk_bytes = src_x * max_src_rows * sizeof(float);
+    const uint32_t dst_chunk_bytes = dst_x * dst_chunk_rows * sizeof(float);
+
+    L->max_src_chunk_rows      = max_src_rows;
+    L->dst_chunk_rows          = dst_chunk_rows;
+    L->src_spad_half_size      = hex_round_up(src_chunk_bytes + 256, 128);
+    L->dst_spad_half_size      = hex_round_up(dst_chunk_bytes, 128);
+    L->src_bytes_per_thread    = 2 * L->src_spad_half_size;
+    L->dst_bytes_per_thread    = 2 * L->dst_spad_half_size;
+    L->r_bytes_per_thread      = hex_round_up(src_x * sizeof(float), 128);
+    L->gather_bytes_per_thread = 256;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_src,    L->src_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_dst,    L->dst_bytes_per_thread * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_r,      L->r_bytes_per_thread   * n_threads);
+    VTCM_LAYOUT_ALLOC(off, off_gather, L->gather_bytes_per_thread * n_threads);
+    L->total_bytes = off;
+}
+
+static inline bool htp_upscale_solve_layout(
+    struct htp_upscale_vtcm_layout * layout,
+    uint32_t src_x,
+    uint32_t src_y,
+    uint32_t dst_x,
+    uint32_t dst_y,
+    float    inv_sf1,
+    float    filter_radius1,
+    uint32_t n_threads,
+    size_t   vtcm_budget,
+    uint32_t * out_dst_chunk_rows
+) {
+    // 1. Try full plane
+    uint32_t chunk_rows = dst_y;
+    htp_upscale_vtcm_layout_build(layout, src_x, src_y, dst_x, dst_y, chunk_rows, inv_sf1, filter_radius1, n_threads);
+    if (layout->total_bytes <= vtcm_budget) {
+        *out_dst_chunk_rows = chunk_rows;
+        return true;
+    }
+
+    // 2. Try smaller chunk bands: 64, 32, 16, 8, 4, 2, 1
+    chunk_rows = 64;
+    while (chunk_rows > 1) {
+        if (chunk_rows < dst_y) {
+            htp_upscale_vtcm_layout_build(layout, src_x, src_y, dst_x, dst_y, chunk_rows, inv_sf1, filter_radius1, n_threads);
+            if (layout->total_bytes <= vtcm_budget) {
+                *out_dst_chunk_rows = chunk_rows;
+                return true;
+            }
+        }
+        chunk_rows /= 2;
+    }
+
+    // 3. Fallback to 1 row
+    chunk_rows = 1;
+    htp_upscale_vtcm_layout_build(layout, src_x, src_y, dst_x, dst_y, chunk_rows, inv_sf1, filter_radius1, n_threads);
+    if (layout->total_bytes <= vtcm_budget) {
+        *out_dst_chunk_rows = chunk_rows;
+        return true;
+    }
+
+    return false;
+}
+
+int op_upscale(struct htp_ops_context * octx);
+
+#endif // HTP_UPSCALE_OPS_H

ggml/src/ggml-impl.h

@@ -512,6 +512,7 @@ static inline float ggml_e8m0_to_fp32_half(uint8_t x) {
 // UE4M3: unsigned, 4 exp bits (bias=7), 3 mantissa bits
 // Returns value * 0.5 to match kvalues_mxfp4 convention (kvalues = 2 * E2M1_float)
 static inline float ggml_ue4m3_to_fp32(uint8_t x) {
+    // TODO: reject invalid unsigned E4M3 scales instead of decoding NaN as zero or ignoring the sign bit.
     if (x == 0 || x == 0x7F) {
         return 0.0f;
     }
@@ -527,6 +528,7 @@ static inline float ggml_ue4m3_to_fp32(uint8_t x) {
 }
 
 static inline uint8_t ggml_fp32_to_ue4m3(float x) {
+    // TODO: use ties-to-even rounding, allow subnormal carry, and preserve finite values with exponent field 15.
     if (!(x > 0.0f)) {
         return 0;
     }
@@ -564,6 +566,62 @@ static inline uint8_t ggml_fp32_to_ue4m3(float x) {
     return (uint8_t) ((ue4m3_exp << 3) | ue4m3_man);
 }
 
+static inline float ggml_f8_e4m3_to_fp32(uint8_t x) {
+    const uint8_t ax = x & 0x7F;
+    if (ax == 0x7F) {
+        // OCP FP8 defines both 0x7F and 0xFF as NaN; the sign need not be preserved.
+        return NAN;
+    }
+
+    const int exp = (ax >> 3) & 0xF;
+    const int man = ax & 0x7;
+    // OCP FP8 rev. 1.1, section 5.1: subnormal magnitude = M * 2^-9; normal magnitude = (1 + M / 8) * 2^(E - 7).
+    const float value = exp == 0
+        ? ldexpf((float) man, -9)
+        : ldexpf(1.0f + (float) man / 8.0f, exp - 7);
+    return x & 0x80 ? -value : value;
+}
+
+static inline int ggml_round_to_nearest_even(float x) {
+    const int value = (int) floorf(x);
+    const float fraction = x - value;
+    return fraction > 0.5f || (fraction == 0.5f && (value & 1)) ? value + 1 : value;
+}
+
+// Supports saturating conversion only; infinities and overflow clamp to +/-448.
+static inline uint8_t ggml_fp32_to_f8_e4m3(float x) {
+    const uint8_t sign = signbit(x) ? 0x80 : 0;
+    x = fabsf(x);
+
+    if (isnan(x)) {
+        // OCP FP8 leaves NaN sign handling to the implementation; use canonical 0x7F.
+        return 0x7F;
+    }
+    if (x == 0.0f) {
+        return sign;
+    }
+    if (isinf(x) || x >= 448.0f) {
+        return sign | 0x7E;
+    }
+    if (x < 0.015625f) {
+        // OCP FP8 rev. 1.1, section 5.1: subnormal magnitude = M * 2^(1 - bias - m) = M / 512.
+        return sign | (uint8_t) ggml_round_to_nearest_even(x * 512.0f);
+    }
+
+    int exp;
+    const float mantissa = frexpf(x, &exp) * 2.0f;
+    int encoded_exp = exp + 6;
+    int encoded_man = ggml_round_to_nearest_even((mantissa - 1.0f) * 8.0f);
+    if (encoded_man == 8) {
+        encoded_man = 0;
+        encoded_exp++;
+    }
+    if (encoded_exp > 15 || (encoded_exp == 15 && encoded_man > 6)) {
+        return sign | 0x7E;
+    }
+    return sign | (uint8_t) (encoded_exp << 3) | (uint8_t) encoded_man;
+}
+
 /**
  * Converts brain16 to float32.
  *

ggml/src/ggml-metal/ggml-metal-device.cpp

@@ -1103,7 +1103,6 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta
             {
                 nsg = N_SG_MXFP4;
                 nr0 = N_R0_MXFP4;
-                smem = 32*sizeof(float);
             } break;
         case GGML_TYPE_Q2_K:
             {
@@ -1399,7 +1398,6 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id(ggml_m
             {
                 nsg = N_SG_MXFP4;
                 nr0 = N_R0_MXFP4;
-                smem = 32*sizeof(float);
             } break;
         case GGML_TYPE_Q2_K:
             {
@@ -1989,6 +1987,23 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext(
     return res;
 }
 
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_rows(
+        ggml_metal_library_t lib,
+        const ggml_tensor * op) {
+    assert(op->op == GGML_OP_FLASH_ATTN_EXT);
+
+    const char * base = "kernel_flash_attn_ext_rows";
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, base);
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, base, base, nullptr);
+    }
+
+    GGML_UNUSED(op);
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec_idx(
         ggml_metal_library_t lib,
         const ggml_tensor * op) {
@@ -2018,6 +2033,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
         bool    has_scap,
         bool    has_kvpad,
         bool    has_sparse,
+        bool    has_rows,
         int32_t nqpsg,
         int32_t ne,
         int32_t nsg,
@@ -2047,7 +2063,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
             dv,
             qne_suffix);
 
-    snprintf(name, 256, "%s_mask=%d_sink=%d_bias=%d_scap=%d_kvpad=%d_sparse=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d",
+    snprintf(name, 256, "%s_mask=%d_sink=%d_bias=%d_scap=%d_kvpad=%d_sparse=%d_rows=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d",
             base,
             has_mask,
             has_sinks,
@@ -2055,6 +2071,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
             has_scap,
             has_kvpad,
             has_sparse,
+            has_rows,
             ns10,
             ns20,
             nsg, nwg);
@@ -2069,6 +2086,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
         ggml_metal_cv_set_bool(cv, has_scap,  FC_FLASH_ATTN_EXT_VEC + 3);
         ggml_metal_cv_set_bool(cv, has_kvpad,  FC_FLASH_ATTN_EXT_VEC + 4);
         ggml_metal_cv_set_bool(cv, has_sparse, FC_FLASH_ATTN_EXT_VEC + 5);
+        ggml_metal_cv_set_bool(cv, has_rows,   FC_FLASH_ATTN_EXT_VEC + 6);
 
         ggml_metal_cv_set_int32(cv, ns10, FC_FLASH_ATTN_EXT_VEC + 20);
         ggml_metal_cv_set_int32(cv, ns20, FC_FLASH_ATTN_EXT_VEC + 21);

ggml/src/ggml-metal/ggml-metal-device.h

@@ -219,6 +219,10 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att
         int32_t ns10,
         int32_t ns20);
 
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_rows(
+        ggml_metal_library_t lib,
+        const struct ggml_tensor * op);
+
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec_idx(
         ggml_metal_library_t lib,
         const struct ggml_tensor * op);
@@ -232,6 +236,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att
         bool    has_scap,
         bool    has_kvpad,
         bool    has_sparse,
+        bool    has_rows,
         int32_t nqpsg,
         int32_t ne,
         int32_t nsg,
@@ -362,7 +367,10 @@ bool   ggml_metal_buffer_is_shared(ggml_metal_buffer_t buf);
 void   ggml_metal_buffer_memset_tensor(ggml_metal_buffer_t buf, struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size);
 void   ggml_metal_buffer_set_tensor   (ggml_metal_buffer_t buf, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size);
 void   ggml_metal_buffer_get_tensor   (ggml_metal_buffer_t buf, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size);
+void   ggml_metal_buffer_set_tensor_n (ggml_metal_buffer_t buf, struct ggml_tensor * const * tensors, const void * const * data, const size_t * offsets, const size_t * sizes, size_t n_tensors);
+void   ggml_metal_buffer_get_tensor_n (ggml_metal_buffer_t buf, const struct ggml_tensor * const * tensors,       void * const * data, const size_t * offsets, const size_t * sizes, size_t n_tensors);
 bool   ggml_metal_buffer_cpy_tensor   (ggml_metal_buffer_t buf, const struct ggml_tensor * src, struct ggml_tensor * dst);
+bool   ggml_metal_buffer_cpy_tensor_n (ggml_metal_buffer_t buf, const struct ggml_tensor * const * srcs, struct ggml_tensor * const * dsts, size_t n_tensors);
 void   ggml_metal_buffer_clear        (ggml_metal_buffer_t buf, uint8_t value);
 
 // finds the Metal buffer that contains the tensor data on the GPU device

ggml/src/ggml-metal/ggml-metal-device.m

@@ -1516,6 +1516,7 @@ static bool ggml_metal_supports_mul_mat_op(
         bool src0_f16_has_mv,
         bool mm_path) {
     if (!has_simdgroup_reduction ||
+        op->src[0]->type == GGML_TYPE_F8_E4M3 ||
         op->src[0]->type == GGML_TYPE_NVFP4 ||
         op->src[0]->type == GGML_TYPE_TQ1_0) {
         return false;
@@ -1726,6 +1727,39 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
         case GGML_OP_ROLL:
             return ggml_is_contiguous(op->src[0]);
         case GGML_OP_FLASH_ATTN_EXT:
+            // K/V row lists (ggml_flash_attn_ext_rows): one list of cache rows per q slice
+            if (op->src[5] != NULL) {
+                const struct ggml_tensor * kv_rows = op->src[5];
+
+                if (kv_rows->type != GGML_TYPE_I32 ||
+                    !ggml_is_contiguous(kv_rows) ||
+                    kv_rows->ne[2] != 1 || kv_rows->ne[3] != 1 ||
+                    kv_rows->ne[1] != op->src[0]->ne[3] ||
+                    !op->src[3] || op->src[1]->ne[3] != 1 || op->src[2]->ne[3] != 1) {
+                    return false;
+                }
+
+                // the non-vec kernels attend over a contiguous F16 copy of the referenced rows:
+                // - F16 K/V is gathered directly
+                // - quantized K/V is first dequantized with kernel_flash_attn_ext_kv_*_f16
+                // - other types keep attending to the whole cache (dense mask)
+                switch (op->src[1]->type) {
+                    case GGML_TYPE_F16:
+                    case GGML_TYPE_Q4_0:
+                    case GGML_TYPE_Q4_1:
+                    case GGML_TYPE_Q5_0:
+                    case GGML_TYPE_Q5_1:
+                    case GGML_TYPE_Q8_0:
+                        break;
+                    default:
+                        return false;
+                }
+
+                // the row gather copies F16 values in groups of 4
+                if (op->src[2]->ne[0] % 4 != 0) {
+                    return false;
+                }
+            }
             // for new head sizes, add checks here
             if (op->src[0]->ne[0] != 32 &&
                 op->src[0]->ne[0] != 40 &&
@@ -1933,7 +1967,8 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
             }
         case GGML_OP_GET_ROWS:
             return op->src[0]->type != GGML_TYPE_NVFP4 &&
-                   op->src[0]->type != GGML_TYPE_TQ1_0;
+                   op->src[0]->type != GGML_TYPE_TQ1_0 &&
+                   op->src[0]->type != GGML_TYPE_F8_E4M3;
         case GGML_OP_SET_ROWS:
             {
                 if (op->src[0]->type == GGML_TYPE_F16) {
@@ -2100,19 +2135,11 @@ static void ggml_metal_buffer_rset_free(ggml_metal_buffer_t buf) {
 static void * ggml_metal_host_malloc(size_t n) {
     void * data = NULL;
 
-#if TARGET_OS_OSX
     kern_return_t err = vm_allocate((vm_map_t) mach_task_self(), (void *) &data, n, VM_FLAGS_ANYWHERE);
     if (err != KERN_SUCCESS) {
         GGML_LOG_ERROR("%s: error: vm_allocate failed\n", __func__);
         return NULL;
     }
-#else
-    const int result = posix_memalign((void **) &data, sysconf(_SC_PAGESIZE), n);
-    if (result != 0) {
-        GGML_LOG_ERROR("%s: error: posix_memalign failed\n", __func__);
-        return NULL;
-    }
-#endif
 
     return data;
 }
@@ -2294,11 +2321,7 @@ void ggml_metal_buffer_free(ggml_metal_buffer_t buf) {
     }
 
     if (buf->is_shared && buf->owned) {
-#if TARGET_OS_OSX
         vm_deallocate((vm_map_t)mach_task_self(), (vm_address_t)buf->all_data, buf->all_size);
-#else
-        free(buf->all_data);
-#endif
     }
 
     free(buf);
@@ -2436,6 +2459,148 @@ void ggml_metal_buffer_get_tensor(ggml_metal_buffer_t buf, const struct ggml_ten
     }
 }
 
+static bool ggml_metal_buffer_data_span(const void * const * data, const size_t * sizes, size_t n_tensors, uintptr_t * min_data, uintptr_t * max_data) {
+    bool found = false;
+
+    for (size_t i = 0; i < n_tensors; i++) {
+        if (sizes[i] == 0) {
+            continue;
+        }
+
+        const uintptr_t start = (uintptr_t) data[i];
+        const uintptr_t end   = start + sizes[i];
+        if (!found) {
+            *min_data = start;
+            *max_data = end;
+            found = true;
+        } else {
+            if (start < *min_data) {
+                *min_data = start;
+            }
+            if (end > *max_data) {
+                *max_data = end;
+            }
+        }
+    }
+
+    return found;
+}
+
+void ggml_metal_buffer_set_tensor_n(ggml_metal_buffer_t buf, struct ggml_tensor * const * tensors, const void * const * data,
+        const size_t * offsets, const size_t * sizes, size_t n_tensors) {
+    if (buf->is_shared) {
+        for (size_t i = 0; i < n_tensors; i++) {
+            if (sizes[i] == 0) {
+                continue;
+            }
+            memcpy((char *) tensors[i]->data + offsets[i], data[i], sizes[i]);
+        }
+        return;
+    }
+
+    uintptr_t min_data = 0, max_data = 0;
+    if (!ggml_metal_buffer_data_span(data, sizes, n_tensors, &min_data, &max_data)) {
+        return;
+    }
+
+    @autoreleasepool {
+        void * data_ptr = (void *) min_data;
+        id<MTLBuffer> buf_src = [buf->dev->mtl_device newBufferWithBytesNoCopy:data_ptr
+                                                               length:max_data - min_data
+                                                              options:MTLResourceStorageModeShared
+                                                          deallocator:nil];
+
+        GGML_ASSERT(buf_src);
+
+        id<MTLCommandBuffer> cmd_buf = [buf->dev->mtl_queue commandBufferWithUnretainedReferences];
+
+        {
+            id<MTLBlitCommandEncoder> encoder = [cmd_buf blitCommandEncoder];
+
+            for (size_t i = 0; i < n_tensors; i++) {
+                if (sizes[i] == 0) {
+                    continue;
+                }
+
+                ggml_backend_buffer_t tensor_buf_ggml = tensors[i]->view_src ? tensors[i]->view_src->buffer : tensors[i]->buffer;
+                ggml_metal_buffer_t tensor_buf = (ggml_metal_buffer_t) tensor_buf_ggml->context;
+                struct ggml_metal_buffer_id bid_dst = ggml_metal_buffer_get_id(tensor_buf, tensors[i]);
+                bid_dst.offs += offsets[i];
+
+                [encoder copyFromBuffer:buf_src
+                           sourceOffset:(uintptr_t) data[i] - min_data
+                               toBuffer:bid_dst.metal
+                      destinationOffset:bid_dst.offs
+                                   size:sizes[i]];
+            }
+
+            [encoder endEncoding];
+        }
+
+        [cmd_buf commit];
+        [cmd_buf waitUntilCompleted];
+
+        [buf_src release];
+    }
+}
+
+void ggml_metal_buffer_get_tensor_n(ggml_metal_buffer_t buf, const struct ggml_tensor * const * tensors, void * const * data,
+        const size_t * offsets, const size_t * sizes, size_t n_tensors) {
+    if (buf->is_shared) {
+        for (size_t i = 0; i < n_tensors; i++) {
+            if (sizes[i] == 0) {
+                continue;
+            }
+            memcpy(data[i], (const char *) tensors[i]->data + offsets[i], sizes[i]);
+        }
+        return;
+    }
+
+    uintptr_t min_data = 0, max_data = 0;
+    if (!ggml_metal_buffer_data_span((const void * const *) data, sizes, n_tensors, &min_data, &max_data)) {
+        return;
+    }
+
+    @autoreleasepool {
+        id<MTLBuffer> buf_dst = [buf->dev->mtl_device newBufferWithBytesNoCopy:(void *) min_data
+                                                               length:max_data - min_data
+                                                              options:MTLResourceStorageModeShared
+                                                          deallocator:nil];
+
+        GGML_ASSERT(buf_dst);
+
+        id<MTLCommandBuffer> cmd_buf = [buf->dev->mtl_queue commandBufferWithUnretainedReferences];
+
+        {
+            id<MTLBlitCommandEncoder> encoder = [cmd_buf blitCommandEncoder];
+
+            for (size_t i = 0; i < n_tensors; i++) {
+                if (sizes[i] == 0) {
+                    continue;
+                }
+
+                ggml_backend_buffer_t tensor_buf_ggml = tensors[i]->view_src ? tensors[i]->view_src->buffer : tensors[i]->buffer;
+                ggml_metal_buffer_t tensor_buf = (ggml_metal_buffer_t) tensor_buf_ggml->context;
+                struct ggml_metal_buffer_id bid_src = ggml_metal_buffer_get_id(tensor_buf, tensors[i]);
+                bid_src.offs += offsets[i];
+
+                [encoder copyFromBuffer:bid_src.metal
+                           sourceOffset:bid_src.offs
+                               toBuffer:buf_dst
+                      destinationOffset:(uintptr_t) data[i] - min_data
+                                   size:sizes[i]];
+            }
+
+            [encoder endEncoding];
+        }
+
+        [cmd_buf commit];
+        [cmd_buf waitUntilCompleted];
+
+        [buf_dst release];
+    }
+}
+
 bool ggml_metal_buffer_cpy_tensor(ggml_metal_buffer_t buf_dst, const struct ggml_tensor * src, struct ggml_tensor * dst) {
     ggml_metal_buffer_t buf_src = (ggml_metal_buffer_t)src->buffer->context;
 
@@ -2477,6 +2642,74 @@ bool ggml_metal_buffer_cpy_tensor(ggml_metal_buffer_t buf_dst, const struct ggml
     return true;
 }
 
+bool ggml_metal_buffer_cpy_tensor_n(ggml_metal_buffer_t buf_dst, const struct ggml_tensor * const * srcs, struct ggml_tensor * const * dsts, size_t n_tensors) {
+    if (n_tensors == 0) {
+        return true;
+    }
+
+    bool all_shared = true;
+    for (size_t i = 0; i < n_tensors; ++i) {
+        ggml_backend_buffer_t src_buf = srcs[i]->view_src ? srcs[i]->view_src->buffer : srcs[i]->buffer;
+        ggml_metal_buffer_t buf_src = (ggml_metal_buffer_t) src_buf->context;
+        if (!buf_dst->is_shared || !buf_src->is_shared) {
+            all_shared = false;
+            break;
+        }
+    }
+
+    if (all_shared) {
+        for (size_t i = 0; i < n_tensors; ++i) {
+            memcpy(dsts[i]->data, srcs[i]->data, ggml_nbytes(srcs[i]));
+        }
+        return true;
+    }
+
+    // for private buffers, we need to use Metal blit commands
+    @autoreleasepool {
+        struct ggml_metal_buffer_id * bid_srcs = malloc(n_tensors * sizeof(struct ggml_metal_buffer_id));
+        struct ggml_metal_buffer_id * bid_dsts = malloc(n_tensors * sizeof(struct ggml_metal_buffer_id));
+
+        for (size_t i = 0; i < n_tensors; ++i) {
+            ggml_backend_buffer_t src_buf = srcs[i]->view_src ? srcs[i]->view_src->buffer : srcs[i]->buffer;
+            ggml_backend_buffer_t dst_buf = dsts[i]->view_src ? dsts[i]->view_src->buffer : dsts[i]->buffer;
+            ggml_metal_buffer_t buf_src = (ggml_metal_buffer_t) src_buf->context;
+            ggml_metal_buffer_t buf_dst_i = (ggml_metal_buffer_t) dst_buf->context;
+            bid_srcs[i] = ggml_metal_buffer_get_id(buf_src, srcs[i]);
+            bid_dsts[i] = ggml_metal_buffer_get_id(buf_dst_i, dsts[i]);
+
+            if (bid_srcs[i].metal == nil || bid_dsts[i].metal == nil) {
+                free(bid_srcs);
+                free(bid_dsts);
+                return false;
+            }
+        }
+
+        id<MTLCommandBuffer> cmd_buf = [buf_dst->dev->mtl_queue commandBufferWithUnretainedReferences];
+
+        {
+            id<MTLBlitCommandEncoder> encoder = [cmd_buf blitCommandEncoder];
+
+            for (size_t i = 0; i < n_tensors; ++i) {
+                [encoder copyFromBuffer:bid_srcs[i].metal
+                           sourceOffset:bid_srcs[i].offs
+                               toBuffer:bid_dsts[i].metal
+                      destinationOffset:bid_dsts[i].offs
+                                   size:ggml_nbytes(srcs[i])];
+            }
+
+            [encoder endEncoding];
+        }
+
+        [cmd_buf commit];
+        [cmd_buf waitUntilCompleted];
+
+        free(bid_srcs);
+        free(bid_dsts);
+    }
+
+    return true;
+}
+
 void ggml_metal_buffer_clear(ggml_metal_buffer_t buf, uint8_t value) {
     if (buf->is_shared) {
         memset(buf->all_data, value, buf->all_size);

ggml/src/ggml-metal/ggml-metal-impl.h

@@ -385,6 +385,19 @@ typedef struct {
     int32_t  nblocks;
 } ggml_metal_kargs_flash_attn_ext_kv_f16;
 
+// gather the K/V cache rows referenced by the kv_rows lists into contiguous F16 buffers
+typedef struct {
+    int32_t  ne11;      // number of rows per list (n_kv_rows)
+    int32_t  ne_12_2;   // number of K/V heads
+    int32_t  ne10;      // DK
+    int32_t  ne20;      // DV, set to 0 to skip gathering V (V is a view of K)
+    uint64_t v_off;     // byte offset of the gathered V region in dst
+    uint64_t nb11;      // K source strides (ne13 = 1: the cache is shared by all streams)
+    uint64_t nb12;
+    uint64_t nb21;      // V source strides (ne23 = 1)
+    uint64_t nb22;
+} ggml_metal_kargs_flash_attn_ext_rows;
+
 typedef struct {
     int32_t  ne11;
     int32_t  ne_12_2; // assume K and V are same shape

ggml/src/ggml-metal/ggml-metal-ops.cpp

@@ -899,7 +899,9 @@ int ggml_metal_op_unary(ggml_metal_op_t ctx, int idx) {
     if (pipeline.cnt) {
         const int n = pipeline.c4 ? ggml_nelements(op)/4 : ggml_nelements(op);
 
-        ggml_metal_encoder_dispatch_threadgroups(enc, n, 1, 1, 1, 1, 1);
+        const int nth = MIN(256, n);
+
+        ggml_metal_encoder_dispatch_threadgroups(enc, (n + nth - 1)/nth, 1, 1, nth, 1, 1);
     } else {
         const int nth_max = MIN(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
         const int nth = MIN(args.ne00, nth_max);
@@ -952,7 +954,9 @@ int ggml_metal_op_glu(ggml_metal_op_t ctx, int idx) {
 
     const int64_t nrows = ggml_nrows(op->src[0]);
 
-    const int32_t nth = std::max(1, std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), ne00/2));
+    // split each row over several threadgroups: with a single row (decode) one threadgroup would run on one GPU core
+    const int32_t nth = std::max(1, std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), std::min(256, ne0)));
+    const int32_t nk  = (ne0 + nth - 1)/nth;
 
     ggml_metal_encoder_set_pipeline(enc, pipeline);
     ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
@@ -964,7 +968,7 @@ int ggml_metal_op_glu(ggml_metal_op_t ctx, int idx) {
     }
     ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(op),         3);
 
-    ggml_metal_encoder_dispatch_threadgroups(enc, nrows, 1, 1, nth, 1, 1);
+    ggml_metal_encoder_dispatch_threadgroups(enc, nk, nrows, 1, nth, 1, 1);
 
     return 1;
 }
@@ -3257,7 +3261,10 @@ static bool ggml_metal_op_flash_attn_ext_use_tensor(const ggml_tensor * op, bool
         return false;
     }
 
-    return op->src[1]->ne[1] % OP_FLASH_ATTN_EXT_TENSOR_NCPSG == 0;
+    // with kv_rows, the attention runs over the compacted row lists (padded to 256), not over the cache
+    const int64_t ne11_attn = op->src[5] ? op->src[5]->ne[0] : op->src[1]->ne[1];
+
+    return ne11_attn % OP_FLASH_ATTN_EXT_TENSOR_NCPSG == 0;
 }
 
 // returns the n_kv_max hint if the sparse path is available for this op, or 0 otherwise
@@ -3266,6 +3273,11 @@ static bool ggml_metal_op_flash_attn_ext_use_tensor(const ggml_tensor * op, bool
 static int ggml_metal_op_flash_attn_ext_n_kv_max_sparse(const ggml_tensor * op) {
     assert(op->op == GGML_OP_FLASH_ATTN_EXT);
 
+    // K/V row lists are handled by the dense kernels, not by the sparse path
+    if (op->src[5] != nullptr) {
+        return 0;
+    }
+
     int32_t n_kv_max = 0;
     memcpy(&n_kv_max, ((const int32_t *) op->op_params) + 4, sizeof(n_kv_max));
 
@@ -3509,6 +3521,32 @@ size_t ggml_metal_op_flash_attn_ext_extra_idx(const ggml_tensor * op) {
     return GGML_PAD(sizeof(int32_t)*(size_t) n_kv_max_padded*ne31*ne32*ne33, 16);
 }
 
+// size of the contiguous F16 buffers holding the K/V rows referenced by the kv_rows lists
+// the attention over these buffers is a standard dense op (no per-tile gather in the main kernel)
+size_t ggml_metal_op_flash_attn_ext_extra_rows(const ggml_tensor * op) {
+    assert(op->op == GGML_OP_FLASH_ATTN_EXT);
+
+    const ggml_tensor * kv_rows = op->src[5];
+
+    if (!kv_rows) {
+        return 0;
+    }
+
+    GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne);
+    GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne);
+
+    const int64_t n_rows   = kv_rows->ne[0];
+    const int64_t n_slices = kv_rows->ne[1];
+
+    // note: when V is a view of K, the gathered V is read from the gathered K, so only K is gathered
+    const bool v_is_view_of_k = ggml_metal_op_flash_attn_ext_v_is_view_of_k(op);
+
+    const size_t k_size = GGML_PAD(sizeof(ggml_fp16_t)*(size_t) ne10*n_rows*ne12*n_slices, 16);
+    const size_t v_size = v_is_view_of_k ? 0 : GGML_PAD(sizeof(ggml_fp16_t)*(size_t) ne20*n_rows*ne22*n_slices, 16);
+
+    return k_size + v_size;
+}
+
 int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
     ggml_tensor * op = ctx->node(idx);
 
@@ -3541,6 +3579,22 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
     GGML_ASSERT(!op->src[3] || op->src[3]->ne[1] >= op->src[0]->ne[1] &&
             "the Flash-Attention Metal kernel requires the mask to be at least n_queries big");
 
+    // K/V row lists (ggml_flash_attn_ext_rows): src[5] holds one list of cache rows per q slice
+    // the attention loops then run over the mask columns and read K/V rows through the lists
+    const ggml_tensor * kv_rows = op->src[5];
+    const bool has_rows = kv_rows != nullptr;
+
+    GGML_ASSERT(!has_rows || (kv_rows->type == GGML_TYPE_I32 && ggml_is_contiguous(kv_rows)));
+    GGML_ASSERT(!has_rows || (kv_rows->ne[1] == ne03 && kv_rows->ne[2] == 1 && kv_rows->ne[3] == 1));
+    GGML_ASSERT(!has_rows || op->src[3]);
+
+    // the loop bound over the mask columns / K/V row lists
+    const int64_t ne11_attn = has_rows ? kv_rows->ne[0] : ne11;
+
+    // the row lists are padded by the KV cache to a multiple of 256, so the pad path never triggers
+    GGML_ASSERT(!has_rows || ne11_attn % OP_FLASH_ATTN_EXT_NCPSG == 0);
+    GGML_ASSERT(!has_rows || ne11_attn % OP_FLASH_ATTN_EXT_VEC_NCPSG == 0);
+
     float scale;
     float max_bias;
     float logit_softcap;
@@ -3555,6 +3609,8 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
 
     const bool has_mask  = op->src[3] != NULL;
     const bool has_sinks = op->src[4] != NULL;
+
+    GGML_ASSERT(!has_rows || has_mask);
     const bool has_bias  = max_bias != 0.0f;
     const bool has_scap  = logit_softcap != 0.0f;
 
@@ -3591,12 +3647,25 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
     const bool use_sparse = n_kv_max_sparse > 0;
     const int n_kv_max_padded = use_sparse ? GGML_PAD(n_kv_max_sparse, OP_FLASH_ATTN_EXT_VEC_NCPSG) : 0;
 
+    const bool use_vec = ggml_metal_op_flash_attn_ext_use_vec(op);
+
+    // with kv_rows, the non-vec path gathers the referenced rows into a contiguous F16 buffer and
+    // attends over it with the F16 kernel - quantized KV is first dequantized with the F16 pass
+    const bool rows_gather = has_rows && !use_vec;
+
     // the vec kernels dequantize the KV inline; no need for the F16 dequant pass in the sparse path
-    const bool use_kv_f16 = !use_sparse && ggml_metal_op_flash_attn_ext_use_kv_f16(op);
+    const bool dequant_kv = !use_sparse && (rows_gather ?
+        ggml_is_quantized(op->src[1]->type) : ggml_metal_op_flash_attn_ext_use_kv_f16(op));
+
+    const bool use_kv_f16 = !use_sparse && (rows_gather ? true : dequant_kv);
 
     ggml_metal_buffer_id bid_idx = bid_kv_f16;
     bid_idx.offs += ggml_metal_op_flash_attn_ext_extra_kv_f16(op);
 
+    // contiguous K/V buffers holding the rows referenced by the kv_rows lists
+    ggml_metal_buffer_id bid_rows = bid_idx;
+    bid_rows.offs += ggml_metal_op_flash_attn_ext_extra_idx(op);
+
     ggml_metal_buffer_id bid_k = bid_src1;
     ggml_metal_buffer_id bid_v = bid_src2;
 
@@ -3609,7 +3678,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
     uint64_t nb22_attn = nb22;
     uint64_t nb23_attn = nb23;
 
-    if (use_kv_f16) {
+    if (dequant_kv) {
         assert(ggml_metal_op_flash_attn_ext_extra_kv_f16(op) != 0);
 
         const bool v_is_view_of_k = ggml_metal_op_flash_attn_ext_v_is_view_of_k(op);
@@ -3697,6 +3766,70 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
         }
     }
 
+    // with kv_rows, gather the cache rows referenced by the lists into contiguous F16 buffers
+    // the attention then runs as a standard dense op over the gathered rows (tensor or half8x8 kernel)
+    if (rows_gather) {
+        assert(ggml_metal_op_flash_attn_ext_extra_rows(op) != 0);
+
+        const bool v_is_view_of_k = ggml_metal_op_flash_attn_ext_v_is_view_of_k(op);
+
+        GGML_ASSERT(ne20 % 4 == 0);
+
+        const size_t k_size = GGML_PAD(sizeof(ggml_fp16_t)*(size_t) ne10*ne11_attn*ne12*ne03, 16);
+
+        ggml_metal_kargs_flash_attn_ext_rows args0 = {
+            /*.ne11    =*/ (int32_t) ne11_attn,
+            /*.ne_12_2 =*/ ne12,
+            /*.ne10    =*/ ne10,
+            /*.ne20    =*/ v_is_view_of_k ? 0 : ne20, // the gathered V is a view of the gathered K
+            /*.v_off   =*/ k_size,
+            /*.nb11    =*/ nb11_attn,
+            /*.nb12    =*/ nb12_attn,
+            /*.nb21    =*/ nb21_attn,
+            /*.nb22    =*/ nb22_attn,
+        };
+
+        auto pipeline0 = ggml_metal_library_get_pipeline_flash_attn_ext_rows(lib, op);
+        const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline0), 64);
+
+        ggml_metal_encoder_set_pipeline(enc, pipeline0);
+        ggml_metal_encoder_set_bytes   (enc, &args0, sizeof(args0), 0);
+        ggml_metal_encoder_set_buffer  (enc, bid_k, 1);
+        ggml_metal_encoder_set_buffer  (enc, bid_v, 2);
+        ggml_metal_encoder_set_buffer  (enc, ggml_metal_get_buffer_id(kv_rows), 3);
+        ggml_metal_encoder_set_buffer  (enc, bid_rows, 4);
+
+        ggml_metal_encoder_dispatch_threadgroups(enc, ne11_attn, ne03, 1, nth, 1, 1);
+
+        ggml_metal_op_concurrency_reset(ctx);
+
+        // contiguous F16 layout of the gathered rows: [DK, n_rows, n_head_kv, n_streams]
+        bid_k = bid_rows;
+
+        nb10_attn = sizeof(ggml_fp16_t);
+        nb11_attn = nb10_attn*ne10;
+        nb12_attn = nb11_attn*ne11_attn;
+        nb13_attn = nb12_attn*ne12;
+
+        // if V is a view of K, the gathered V is read from the gathered K with K's strides
+        if (v_is_view_of_k) {
+            bid_v = bid_k;
+
+            nb20_attn = nb10_attn;
+            nb21_attn = nb11_attn;
+            nb22_attn = nb12_attn;
+            nb23_attn = nb13_attn;
+        } else {
+            bid_v = bid_rows;
+            bid_v.offs += k_size;
+
+            nb20_attn = sizeof(ggml_fp16_t);
+            nb21_attn = nb20_attn*ne20;
+            nb22_attn = nb21_attn*ne11_attn;
+            nb23_attn = nb22_attn*ne22;
+        }
+    }
+
     if (!use_sparse && ggml_metal_op_flash_attn_ext_use_tensor(op, props_dev->has_tensor)) {
         // tensor API kernel
         const int nqptg = ne00 >= 512 ? OP_FLASH_ATTN_EXT_TENSOR_NQPSG_LARGE : OP_FLASH_ATTN_EXT_TENSOR_NQPSG; // queries per threadgroup
@@ -3742,9 +3875,9 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
             /*.nb01          =*/ nb01,
             /*.nb02          =*/ nb02,
             /*.nb03          =*/ nb03,
-            /*.ne11          =*/ ne11,
+            /*.ne11          =*/ (int32_t) ne11_attn,
             /*.ne_12_2       =*/ ne12,
-            /*.ne_12_3       =*/ ne13,
+            /*.ne_12_3       =*/ has_rows ? ne03 : ne13, // the gathered K/V holds one copy per stream
             /*.ns10          =*/ ns10,
             /*.nb11          =*/ nb11_attn,
             /*.nb12          =*/ nb12_attn,
@@ -3802,7 +3935,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
 
         bool need_sync = false;
 
-        const bool has_kvpad = ne11 % ncpsg != 0;
+        const bool has_kvpad = ne11_attn % ncpsg != 0;
 
         if (has_kvpad) {
             assert(ggml_metal_op_flash_attn_ext_extra_pad(op) != 0);
@@ -3902,9 +4035,9 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
             /*.nb01          =*/ nb01,
             /*.nb02          =*/ nb02,
             /*.nb03          =*/ nb03,
-            /*.ne11          =*/ ne11,
+            /*.ne11          =*/ (int32_t) ne11_attn,
             /*.ne_12_2       =*/ ne12,
-            /*.ne_12_3       =*/ ne13,
+            /*.ne_12_3       =*/ has_rows ? ne03 : ne13, // the gathered K/V holds one copy per stream
             /*.ns10          =*/ ns10,
             /*.nb11          =*/ nb11_attn,
             /*.nb12          =*/ nb12_attn,
@@ -3970,7 +4103,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
 
         bool need_sync = false;
 
-        const bool has_kvpad = !use_sparse && ne11 % ncpsg != 0;
+        const bool has_kvpad = !use_sparse && ne11_attn % ncpsg != 0;
 
         if (use_sparse) {
             assert(ggml_metal_op_flash_attn_ext_extra_idx(op) != 0);
@@ -4053,13 +4186,13 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
         //   queries:      Q*NSG*PAD2(ne00, 128)
         //   attn + mask:  NSG*4*Q*C
         //   results:      2*NSG*Q*PAD2(ne20, 128)
-        //   sparse idx:   NSG*C ints (only when use_sparse)
+        //   sparse idx:   NSG*C ints (when use_sparse, or when gathering K/V rows)
         const int64_t dk_pad = GGML_PAD(ne00, 128);
         const int64_t dv_pad = GGML_PAD(ne20, 128);
 
         auto fa_vec_smem = [&](int64_t nsg, int32_t nqptg) -> size_t {
             const size_t smem_half = (size_t) (dk_pad + 4*ncpsg + 2*dv_pad)*nqptg*nsg;
-            return GGML_PAD(smem_half*sizeof(ggml_fp16_t) + (use_sparse ? (size_t) nsg*ncpsg*sizeof(int) : 0), 16);
+            return GGML_PAD(smem_half*sizeof(ggml_fp16_t) + ((use_sparse || has_rows) ? (size_t) nsg*ncpsg*sizeof(int) : 0), 16);
         };
 
         int64_t nsg = 1;
@@ -4090,7 +4223,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
         } else {
             nwg = 32;
             nsg = 1;
-            while (2*nwg*nsg*ncpsg < ne11 && nsg < 4) {
+            while (2*nwg*nsg*ncpsg < ne11_attn && nsg < 4) {
                 nsg *= 2;
             }
         }
@@ -4111,7 +4244,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
             /*.nb01          =*/ nb01,
             /*.nb02          =*/ nb02,
             /*.nb03          =*/ nb03,
-            /*.ne11          =*/ use_sparse ? n_kv_max_padded : ne11,
+            /*.ne11          =*/ use_sparse ? n_kv_max_padded : (int32_t) ne11_attn,
             /*.ne_12_2       =*/ ne12,
             /*.ne_12_3       =*/ ne13,
             /*.ns10          =*/ ns10,
@@ -4140,7 +4273,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
             /*.n_kv_max_padded =*/ n_kv_max_padded,
         };
 
-        auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, use_sparse, nqptg, cfg.NE, nsg, nwg, use_kv_f16, ns10, ns20);
+        auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, use_sparse, has_rows, nqptg, cfg.NE, nsg, nwg, use_kv_f16, ns10, ns20);
 
         GGML_ASSERT(nsg*32 <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
 
@@ -4151,7 +4284,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
         ggml_metal_encoder_set_buffer  (enc, bid_v,    3);
         ggml_metal_encoder_set_buffer  (enc, bid_src3, 4);
         ggml_metal_encoder_set_buffer  (enc, bid_src4, 5);
-        ggml_metal_encoder_set_buffer  (enc, use_sparse ? bid_idx : bid_src0, 8);
+        ggml_metal_encoder_set_buffer  (enc, use_sparse ? bid_idx : (has_rows ? ggml_metal_get_buffer_id(kv_rows) : bid_src0), 8);
 
         const size_t smem = fa_vec_smem(nsg, nqptg);
 
@@ -4328,7 +4461,9 @@ int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) {
     ggml_metal_encoder_set_buffer  (enc, bid_dst,  3);
 
     if (pipeline.cnt) {
-        ggml_metal_encoder_dispatch_threadgroups(enc, args.ne0, ggml_nrows(op), 1, 1, 1, 1);
+        const int nth = MIN(256, args.ne0);
+
+        ggml_metal_encoder_dispatch_threadgroups(enc, (args.ne0 + nth - 1)/nth, ggml_nrows(op), 1, nth, 1, 1);
     } else {
         const int nth_max = MIN(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
 

ggml/src/ggml-metal/ggml-metal-ops.h

@@ -45,6 +45,7 @@ size_t ggml_metal_op_flash_attn_ext_extra_blk(const struct ggml_tensor * op);
 size_t ggml_metal_op_flash_attn_ext_extra_tmp(const struct ggml_tensor * op);
 size_t ggml_metal_op_flash_attn_ext_extra_kv_f16(const struct ggml_tensor * op);
 size_t ggml_metal_op_flash_attn_ext_extra_idx(const struct ggml_tensor * op);
+size_t ggml_metal_op_flash_attn_ext_extra_rows(const struct ggml_tensor * op);
 
 size_t ggml_metal_op_argmax_extra(const struct ggml_tensor * op);
 

ggml/src/ggml-metal/ggml-metal.cpp

@@ -68,6 +68,22 @@ static void ggml_backend_metal_buffer_shared_get_tensor(ggml_backend_buffer_t bu
     ggml_metal_buffer_get_tensor(ctx, tensor, data, offset, size);
 }
 
+static void ggml_backend_metal_buffer_shared_set_tensor_n(ggml_backend_buffer_t buffer, struct ggml_tensor * const * tensors, const void * const * data, const size_t * offsets, const size_t * sizes, size_t n_tensors) {
+    ggml_metal_buffer_t ctx = (ggml_metal_buffer_t)buffer->context;
+
+    GGML_ASSERT(ggml_metal_buffer_is_shared(ctx));
+
+    ggml_metal_buffer_set_tensor_n(ctx, tensors, data, offsets, sizes, n_tensors);
+}
+
+static void ggml_backend_metal_buffer_shared_get_tensor_n(ggml_backend_buffer_t buffer, const struct ggml_tensor * const * tensors, void * const * data, const size_t * offsets, const size_t * sizes, size_t n_tensors) {
+    ggml_metal_buffer_t ctx = (ggml_metal_buffer_t)buffer->context;
+
+    GGML_ASSERT(ggml_metal_buffer_is_shared(ctx));
+
+    ggml_metal_buffer_get_tensor_n(ctx, tensors, data, offsets, sizes, n_tensors);
+}
+
 static bool ggml_backend_metal_buffer_shared_cpy_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * src, ggml_tensor * dst) {
     ggml_metal_buffer_t ctx = (ggml_metal_buffer_t)buffer->context;
 
@@ -80,6 +96,20 @@ static bool ggml_backend_metal_buffer_shared_cpy_tensor(ggml_backend_buffer_t bu
     return ggml_metal_buffer_cpy_tensor(ctx, src, dst);
 }
 
+static bool ggml_backend_metal_buffer_shared_cpy_tensor_n(ggml_backend_buffer_t buffer, const struct ggml_tensor * const * srcs, struct ggml_tensor * const * dsts, size_t n_tensors) {
+    ggml_metal_buffer_t ctx = (ggml_metal_buffer_t)buffer->context;
+
+    GGML_ASSERT(ggml_metal_buffer_is_shared(ctx));
+
+    for (size_t i = 0; i < n_tensors; ++i) {
+        if (!ggml_backend_buffer_is_metal(srcs[i]->buffer)) {
+            return false;
+        }
+    }
+
+    return ggml_metal_buffer_cpy_tensor_n(ctx, srcs, dsts, n_tensors);
+}
+
 static void ggml_backend_metal_buffer_shared_clear(ggml_backend_buffer_t buffer, uint8_t value) {
     ggml_metal_buffer_t ctx = (ggml_metal_buffer_t)buffer->context;
 
@@ -100,6 +130,9 @@ static ggml_backend_buffer_i ggml_backend_metal_buffer_shared_i = {
     /* .cpy_tensor    = */ ggml_backend_metal_buffer_shared_cpy_tensor,
     /* .clear         = */ ggml_backend_metal_buffer_shared_clear,
     /* .reset         = */ NULL,
+    /* .set_tensor_n  = */ ggml_backend_metal_buffer_shared_set_tensor_n,
+    /* .get_tensor_n  = */ ggml_backend_metal_buffer_shared_get_tensor_n,
+    /* .cpy_tensor_n  = */ ggml_backend_metal_buffer_shared_cpy_tensor_n,
 };
 
 // private buffer
@@ -144,6 +177,22 @@ static void ggml_backend_metal_buffer_private_get_tensor(ggml_backend_buffer_t b
     ggml_metal_buffer_get_tensor(ctx, tensor, data, offset, size);
 }
 
+static void ggml_backend_metal_buffer_private_set_tensor_n(ggml_backend_buffer_t buffer, struct ggml_tensor * const * tensors, const void * const * data, const size_t * offsets, const size_t * sizes, size_t n_tensors) {
+    ggml_metal_buffer_t ctx = (ggml_metal_buffer_t)buffer->context;
+
+    GGML_ASSERT(!ggml_metal_buffer_is_shared(ctx));
+
+    ggml_metal_buffer_set_tensor_n(ctx, tensors, data, offsets, sizes, n_tensors);
+}
+
+static void ggml_backend_metal_buffer_private_get_tensor_n(ggml_backend_buffer_t buffer, const struct ggml_tensor * const * tensors, void * const * data, const size_t * offsets, const size_t * sizes, size_t n_tensors) {
+    ggml_metal_buffer_t ctx = (ggml_metal_buffer_t)buffer->context;
+
+    GGML_ASSERT(!ggml_metal_buffer_is_shared(ctx));
+
+    ggml_metal_buffer_get_tensor_n(ctx, tensors, data, offsets, sizes, n_tensors);
+}
+
 static bool ggml_backend_metal_buffer_private_cpy_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * src, ggml_tensor * dst) {
     ggml_metal_buffer_t ctx = (ggml_metal_buffer_t)buffer->context;
 
@@ -156,6 +205,20 @@ static bool ggml_backend_metal_buffer_private_cpy_tensor(ggml_backend_buffer_t b
     return ggml_metal_buffer_cpy_tensor(ctx, src, dst);
 }
 
+static bool ggml_backend_metal_buffer_private_cpy_tensor_n(ggml_backend_buffer_t buffer, const struct ggml_tensor * const * srcs, struct ggml_tensor * const * dsts, size_t n_tensors) {
+    ggml_metal_buffer_t ctx = (ggml_metal_buffer_t)buffer->context;
+
+    GGML_ASSERT(!ggml_metal_buffer_is_shared(ctx));
+
+    for (size_t i = 0; i < n_tensors; ++i) {
+        if (!ggml_backend_buffer_is_metal(srcs[i]->buffer)) {
+            return false;
+        }
+    }
+
+    return ggml_metal_buffer_cpy_tensor_n(ctx, srcs, dsts, n_tensors);
+}
+
 static void ggml_backend_metal_buffer_private_clear(ggml_backend_buffer_t buffer, uint8_t value) {
     ggml_metal_buffer_t ctx = (ggml_metal_buffer_t)buffer->context;
 
@@ -176,6 +239,9 @@ static ggml_backend_buffer_i ggml_backend_metal_buffer_private_i = {
     /* .cpy_tensor    = */ ggml_backend_metal_buffer_private_cpy_tensor,
     /* .clear         = */ ggml_backend_metal_buffer_private_clear,
     /* .reset         = */ NULL,
+    /* .set_tensor_n  = */ ggml_backend_metal_buffer_private_set_tensor_n,
+    /* .get_tensor_n  = */ ggml_backend_metal_buffer_private_get_tensor_n,
+    /* .cpy_tensor_n  = */ ggml_backend_metal_buffer_private_cpy_tensor_n,
 };
 
 static bool ggml_backend_buffer_is_metal(ggml_backend_buffer_t buffer) {
@@ -234,6 +300,7 @@ static size_t ggml_backend_metal_buffer_type_get_alloc_size(ggml_backend_buffer_
                 res += ggml_metal_op_flash_attn_ext_extra_tmp(tensor);
                 res += ggml_metal_op_flash_attn_ext_extra_kv_f16(tensor);
                 res += ggml_metal_op_flash_attn_ext_extra_idx(tensor);
+                res += ggml_metal_op_flash_attn_ext_extra_rows(tensor);
             } break;
         case GGML_OP_CUMSUM:
         case GGML_OP_ARGSORT:

ggml/src/ggml-metal/kernels/binbcast.metal

@@ -22,8 +22,13 @@ kernel void kernel_bin_fuse_impl(
 
     if (FC_RB) {
         // row broadcast
-        const uint i0 = tgpig.y*args.ne00 + tgpig.x;
-        const uint i1 = FC_CB ? tgpig.x%args.ne10 : tgpig.x;
+        const uint i = tgpig.x*ntg.x + tpitg.x;
+        if (i >= (uint) args.ne0) {
+            return;
+        }
+
+        const uint i0 = tgpig.y*args.ne00 + i;
+        const uint i1 = FC_CB ? i%args.ne10 : i;
 
         device const T0 * src0_row = (device const T0 *) (src0);
         device       T  * dst_row  = (device       T  *) (dst);