diff --git a/.devops/openvino.Dockerfile b/.devops/openvino.Dockerfile index a43e5c4993f..13301ba287d 100644 --- a/.devops/openvino.Dockerfile +++ b/.devops/openvino.Dockerfile @@ -1,12 +1,12 @@ -ARG OPENVINO_VERSION_MAJOR=2026.3 -ARG OPENVINO_VERSION_FULL=2026.3.0.22451.bd8d6542e3c +ARG OPENVINO_VERSION_MAJOR=2026.3.1 +ARG OPENVINO_VERSION_FULL=2026.3.1.22476.56d9685302d ARG UBUNTU_VERSION=24.04 # Intel GPU driver versions. https://github.com/intel/compute-runtime/releases -ARG IGC_VERSION=v2.38.2 -ARG IGC_VERSION_FULL=2_2.38.2+22051 -ARG COMPUTE_RUNTIME_VERSION=26.27.39122.11 -ARG COMPUTE_RUNTIME_VERSION_FULL=26.27.39122.11-0 +ARG IGC_VERSION=v2.40.13 +ARG IGC_VERSION_FULL=2_2.40.13+22418 +ARG COMPUTE_RUNTIME_VERSION=26.31.39395.13 +ARG COMPUTE_RUNTIME_VERSION_FULL=26.31.39395.13-0 ARG IGDGMM_VERSION=22.10.0 # Intel NPU driver versions. https://github.com/intel/linux-npu-driver/releases diff --git a/.github/actions/ccache-buckets/action.yml b/.github/actions/ccache-buckets/action.yml index 8eb65d275cb..eaa8d164ebf 100644 --- a/.github/actions/ccache-buckets/action.yml +++ b/.github/actions/ccache-buckets/action.yml @@ -58,34 +58,38 @@ runs: if: ${{ inputs.save == 'true' }} shell: bash run: | - set +e -uo pipefail - source .venv-hf/bin/activate - CCACHE_DIR=$(ccache -k cache_dir) - if [[ -d "$CCACHE_DIR" ]]; then - ccache -s - if [[ -n "${{ inputs.evict-old-files }}" ]]; then - ccache --evict-older-than "${{ inputs.evict-old-files }}" - fi - DATESTAMP=$(date -u +'%Y-%m-%dT%H:%M:%SZ') - CACHEFILE="${{ inputs.key }}-$DATESTAMP.tar.gz" - if tar -czf ccache_bucket.tar.gz -C "$CCACHE_DIR" .; then - hf buckets cp ccache_bucket.tar.gz "hf://buckets/${{ inputs.hf_bucket }}/${{ inputs.folder }}/$CACHEFILE" + if [[ -n "$HF_TOKEN" ]]; then + set +e -uo pipefail + source .venv-hf/bin/activate + CCACHE_DIR=$(ccache -k cache_dir) + if [[ -d "$CCACHE_DIR" ]]; then + ccache -s + if [[ -n "${{ inputs.evict-old-files }}" ]]; then + ccache --evict-older-than "${{ inputs.evict-old-files }}" + fi + DATESTAMP=$(date -u +'%Y-%m-%dT%H:%M:%SZ') + CACHEFILE="${{ inputs.key }}-$DATESTAMP.tar.gz" + if tar -czf ccache_bucket.tar.gz -C "$CCACHE_DIR" .; then + hf buckets cp ccache_bucket.tar.gz "hf://buckets/${{ inputs.hf_bucket }}/${{ inputs.folder }}/$CACHEFILE" + fi + rm ccache_bucket.tar.gz + else + echo "'$CCACHE_DIR' not found." fi - rm ccache_bucket.tar.gz - else - echo "'$CCACHE_DIR' not found." fi - name: Remove old ccache files from buckets if: ${{ inputs.save == 'true' }} shell: bash run: | - set +e -uo pipefail - source .venv-hf/bin/activate - CACHE_FILES=$(hf buckets list "hf://buckets/${{ inputs.hf_bucket }}/${{ inputs.folder }}" --json | jq -r '[.[] | select(.type == "file") | select((.uploaded_at | .[:19]+"Z" | fromdateiso8601) < (now - 5 * 60)) | select(.path | startswith("${{ inputs.folder }}/${{ inputs.key }}") and endswith(".tar.gz"))] | sort_by(.path)[:-1] | .[] | [.path // ""] | @tsv') - if [[ -n "$CACHE_FILES" ]]; then - echo "Removing old ccache files..." - while IFS=$'\t' read -r CACHE_PATH; do - hf buckets rm "hf://buckets/${{ inputs.hf_bucket }}/$CACHE_PATH" -y - done <<< "$CACHE_FILES" + if [[ -n "$HF_TOKEN" ]]; then + set +e -uo pipefail + source .venv-hf/bin/activate + CACHE_FILES=$(hf buckets list "hf://buckets/${{ inputs.hf_bucket }}/${{ inputs.folder }}" --json | jq -r '[.[] | select(.type == "file") | select((.uploaded_at | .[:19]+"Z" | fromdateiso8601) < (now - 5 * 60)) | select(.path | startswith("${{ inputs.folder }}/${{ inputs.key }}") and endswith(".tar.gz"))] | sort_by(.path)[:-1] | .[] | [.path // ""] | @tsv') + if [[ -n "$CACHE_FILES" ]]; then + echo "Removing old ccache files..." + while IFS=$'\t' read -r CACHE_PATH; do + hf buckets rm "hf://buckets/${{ inputs.hf_bucket }}/$CACHE_PATH" -y + done <<< "$CACHE_FILES" + fi fi diff --git a/.github/workflows/build-cache.yml b/.github/workflows/build-cache.yml index 187427a8d4b..4a23ec2d4d3 100644 --- a/.github/workflows/build-cache.yml +++ b/.github/workflows/build-cache.yml @@ -41,8 +41,8 @@ jobs: env: # Sync versions in build-openvino.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile - OPENVINO_VERSION_MAJOR: "2026.3" - OPENVINO_VERSION_FULL: "2026.3.0.22451.bd8d6542e3c" + OPENVINO_VERSION_MAJOR: "2026.3.1" + OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d" steps: - name: Clone @@ -69,8 +69,8 @@ jobs: env: # Sync versions in build.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile - OPENVINO_VERSION_MAJOR: "2026.3" - OPENVINO_VERSION_FULL: "2026.3.0.22451.bd8d6542e3c" + OPENVINO_VERSION_MAJOR: "2026.3.1" + OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d" steps: - name: Clone diff --git a/.github/workflows/build-cuda-ubuntu.yml b/.github/workflows/build-cuda-ubuntu.yml index 215299670ac..80bd7820955 100644 --- a/.github/workflows/build-cuda-ubuntu.yml +++ b/.github/workflows/build-cuda-ubuntu.yml @@ -65,7 +65,7 @@ jobs: with: key: cuda-ubuntu-24.04-cuda folder: llama.cpp - hf_bucket: ${{ vars.HF_BUCKET_CACHE_OUTPUT }} + hf_bucket: ggml-org/cache - name: Build with CMake # TODO: Remove GGML_CUDA_CUB_3DOT2 flag once CCCL 3.2 is bundled within CTK and that CTK version is used in this project @@ -89,7 +89,7 @@ jobs: key: cuda-ubuntu-24.04-cuda folder: llama.cpp evict-old-files: 1d - hf_bucket: ${{ vars.HF_BUCKET_CACHE_OUTPUT }} + hf_bucket: ggml-org/cache save: true hip: @@ -120,7 +120,7 @@ jobs: with: key: cuda-ubuntu-22.04-hip folder: llama.cpp - hf_bucket: ${{ vars.HF_BUCKET_CACHE_OUTPUT }} + hf_bucket: ggml-org/cache - name: Build with native CMake HIP support id: cmake_build @@ -140,7 +140,7 @@ jobs: key: cuda-ubuntu-22.04-hip folder: llama.cpp evict-old-files: 1d - hf_bucket: ${{ vars.HF_BUCKET_CACHE_OUTPUT }} + hf_bucket: ggml-org/cache save: true musa: @@ -171,7 +171,7 @@ jobs: with: key: cuda-ubuntu-22.04-musa folder: llama.cpp - hf_bucket: ${{ vars.HF_BUCKET_CACHE_OUTPUT }} + hf_bucket: ggml-org/cache - name: Build with native CMake MUSA support id: cmake_build @@ -189,5 +189,5 @@ jobs: key: cuda-ubuntu-22.04-musa folder: llama.cpp evict-old-files: 1d - hf_bucket: ${{ vars.HF_BUCKET_CACHE_OUTPUT }} + hf_bucket: ggml-org/cache save: true diff --git a/.github/workflows/build-openvino.yml b/.github/workflows/build-openvino.yml index 0316e7ad97e..8e0326f4a44 100644 --- a/.github/workflows/build-openvino.yml +++ b/.github/workflows/build-openvino.yml @@ -32,6 +32,8 @@ env: LLAMA_ARG_LOG_COLORS: 1 LLAMA_ARG_LOG_PREFIX: 1 LLAMA_ARG_LOG_TIMESTAMPS: 1 + # TODO: fix and re-enable the `test-llama-archs` and `test-recurrent-state-rollback` + CTEST_EXCLUDE: "test-llama-archs|^test-recurrent-state-rollback" jobs: ubuntu-24-openvino: @@ -39,8 +41,8 @@ jobs: env: # Sync versions in build-openvino.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile - OPENVINO_VERSION_MAJOR: "2026.3" - OPENVINO_VERSION_FULL: "2026.3.0.22451.bd8d6542e3c" + OPENVINO_VERSION_MAJOR: "2026.3.1" + OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d" steps: - name: Clone @@ -78,26 +80,24 @@ jobs: - name: Test (CPU) id: cmake_test_cpu - # TODO: fix and re-enable the `test-llama-archs` test below run: | cd ${{ github.workspace }} - ctest --test-dir build/ReleaseOV -L main -E "test-llama-archs|test-recurrent-state-rollback-nemotron-h" --verbose --timeout 2000 + ctest --test-dir build/ReleaseOV -L main -E "${{ env.CTEST_EXCLUDE }}" --verbose --timeout 3000 - name: Test (GPU) id: cmake_test_gpu - # TODO: fix and re-enable the `test-llama-archs` test below run: | cd ${{ github.workspace }} export GGML_OPENVINO_DEVICE=GPU - ctest --test-dir build/ReleaseOV -L main -E "test-llama-archs|test-recurrent-state-rollback-nemotron-h" --verbose --timeout 3000 + ctest --test-dir build/ReleaseOV -L main -E "${{ env.CTEST_EXCLUDE }}" --verbose --timeout 3000 openvino-windows-2022: runs-on: windows-2022 env: # Sync versions in build-openvino.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile - OPENVINO_VERSION_MAJOR: "2026.3" - OPENVINO_VERSION_FULL: "2026.3.0.22451.bd8d6542e3c" + OPENVINO_VERSION_MAJOR: "2026.3.1" + OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d" steps: - name: Clone @@ -159,14 +159,13 @@ jobs: - name: Test (CPU) id: cmake_test_cpu shell: cmd - # TODO: fix and re-enable the `test-llama-archs` test below run: | REM Find extracted OpenVINO folder dynamically for /d %%i in (openvino_toolkit\*) do set OPENVINO_ROOT=%%i call "%OPENVINO_ROOT%\setupvars.bat" cd build - ctest --test-dir ReleaseOV -L main -E "test-llama-archs|test-recurrent-state-rollback-nemotron-h" -C Release --verbose --timeout 3000 + ctest --test-dir ReleaseOV -L main -E "${{ env.CTEST_EXCLUDE }}" -C Release --verbose --timeout 3000 - name: ccache-clear uses: ./.github/actions/ccache-clear diff --git a/.github/workflows/build-self-hosted.yml b/.github/workflows/build-self-hosted.yml index fe2ab815473..ccfe2a60464 100644 --- a/.github/workflows/build-self-hosted.yml +++ b/.github/workflows/build-self-hosted.yml @@ -288,8 +288,8 @@ jobs: env: # Sync versions in build.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile - OPENVINO_VERSION_MAJOR: "2026.3" - OPENVINO_VERSION_FULL: "2026.3.0.22451.bd8d6542e3c" + OPENVINO_VERSION_MAJOR: "2026.3.1" + OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d" steps: - name: Clone diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml index be2f840b1af..1de25b522de 100644 --- a/.github/workflows/docker.yml +++ b/.github/workflows/docker.yml @@ -64,7 +64,7 @@ jobs: needs: create_tag uses: ./.github/workflows/ui-build.yml with: - hf_ui_version: ${{ needs.create_tag.outputs.source_tag }} + ui_version: ${{ needs.create_tag.outputs.source_tag }} prepare_matrices: name: Prepare Docker matrices @@ -162,7 +162,7 @@ jobs: if: ${{ matrix.config.prebuilt_ui == true }} uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8 with: - name: ui-build + name: llama-ui.zip path: tools/ui/dist - name: Set up QEMU diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 5815e5c4759..76717d064bd 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -61,31 +61,8 @@ jobs: echo "should_release=false" >> $GITHUB_OUTPUT fi - get-version: - runs-on: ubuntu-slim - outputs: - ui_version: ${{ steps.version.outputs.ui_version }} - steps: - - uses: actions/checkout@v6 - with: - fetch-depth: 0 - - id: version - run: | - # Resolve UI version: BUILD_NUMBER from cmake/build-info.cmake > git hash + epoch > fallback - version="" - if grep -q "BUILD_NUMBER" cmake/build-info.cmake; then - build_number=$(grep "set(BUILD_NUMBER" cmake/build-info.cmake | grep -oP '\d+') - if [ -n "$build_number" ] && [ "$build_number" -gt 0 ]; then - version="b${build_number}" - fi - fi - if [ -z "$version" ]; then - version=$(git rev-parse --short HEAD)-$(date +%s) - fi - echo "ui_version=${version}" >> $GITHUB_OUTPUT - macos-cpu: - needs: [check-release, get-version] + needs: [check-release, ui-build] if: ${{ needs.check-release.outputs.should_release == 'true' }} strategy: matrix: @@ -119,12 +96,11 @@ jobs: with: fetch-depth: 0 - - name: Setup Node.js - uses: actions/setup-node@v6 + - name: Download UI build + uses: actions/download-artifact@v7 with: - node-version: "24" - cache: "npm" - cache-dependency-path: "tools/ui/package-lock.json" + name: llama-ui.zip + path: tools/ui/dist - name: ccache uses: ggml-org/ccache-action@v1.2.21 @@ -141,7 +117,6 @@ jobs: -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \ -DLLAMA_FATAL_WARNINGS=ON \ -DLLAMA_BUILD_BORINGSSL=ON \ - -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \ ${{ env.CMAKE_ARGS }} cmake --build build --config Release -j $(sysctl -n hw.logicalcpu) @@ -167,7 +142,7 @@ jobs: key: release-${{ matrix.os }}-${{ matrix.arch }} ubuntu-cpu: - needs: [check-release, get-version] + needs: [check-release, ui-build] if: ${{ needs.check-release.outputs.should_release == 'true' }} strategy: matrix: @@ -191,12 +166,11 @@ jobs: with: fetch-depth: 0 - - name: Setup Node.js - uses: actions/setup-node@v6 + - name: Download UI build + uses: actions/download-artifact@v7 with: - node-version: "24" - cache: "npm" - cache-dependency-path: "tools/ui/package-lock.json" + name: llama-ui.zip + path: tools/ui/dist - name: Dependencies id: depends @@ -227,7 +201,6 @@ jobs: -DGGML_NATIVE=OFF \ -DGGML_CPU_ALL_VARIANTS=ON \ -DLLAMA_FATAL_WARNINGS=ON \ - -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \ ${{ env.CMAKE_ARGS }} cmake --build build --config Release -j $(nproc) @@ -254,7 +227,7 @@ jobs: key: release-${{ matrix.os }}-cpu ubuntu-vulkan: - needs: [check-release, get-version] + needs: [check-release, ui-build] if: ${{ needs.check-release.outputs.should_release == 'true' }} strategy: @@ -277,12 +250,11 @@ jobs: with: fetch-depth: 0 - - name: Setup Node.js - uses: actions/setup-node@v6 + - name: Download UI build + uses: actions/download-artifact@v7 with: - node-version: "24" - cache: "npm" - cache-dependency-path: "tools/ui/package-lock.json" + name: llama-ui.zip + path: tools/ui/dist - name: Dependencies id: depends @@ -314,7 +286,6 @@ jobs: -DGGML_NATIVE=OFF \ -DGGML_CPU_ALL_VARIANTS=ON \ -DGGML_VULKAN=ON \ - -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \ ${{ env.CMAKE_ARGS }} cmake --build build --config Release -j $(nproc) @@ -340,7 +311,7 @@ jobs: key: release-${{ matrix.os }}-vulkan android-arm64: - needs: [check-release, get-version] + needs: [check-release, ui-build] if: ${{ needs.check-release.outputs.should_release == 'true' }} runs-on: ubuntu-latest @@ -358,12 +329,11 @@ jobs: with: fetch-depth: 0 - - name: Setup Node.js - uses: actions/setup-node@v6 + - name: Download UI build + uses: actions/download-artifact@v7 with: - node-version: "24" - cache: "npm" - cache-dependency-path: "tools/ui/package-lock.json" + name: llama-ui.zip + path: tools/ui/dist - name: Set up JDK uses: actions/setup-java@v5 @@ -407,7 +377,6 @@ jobs: -DLLAMA_FATAL_WARNINGS=ON \ -DGGML_OPENMP=OFF \ -DLLAMA_BUILD_BORINGSSL=ON \ - -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \ ${{ env.CMAKE_ARGS }} cmake --build build --config Release -j $(nproc) @@ -433,7 +402,7 @@ jobs: name: llama-bin-android-arm64.tar.gz ubuntu-24-openvino: - needs: [check-release, get-version] + needs: [check-release, ui-build] if: ${{ needs.check-release.outputs.should_release == 'true' }} runs-on: ubuntu-24.04 @@ -446,8 +415,8 @@ jobs: env: # Sync versions in build-openvino.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile - OPENVINO_VERSION_MAJOR: "2026.3" - OPENVINO_VERSION_FULL: "2026.3.0.22451.bd8d6542e3c" + OPENVINO_VERSION_MAJOR: "2026.3.1" + OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d" steps: - name: Set OpenVINO version output @@ -460,12 +429,11 @@ jobs: with: fetch-depth: 0 - - name: Setup Node.js - uses: actions/setup-node@v6 + - name: Download UI build + uses: actions/download-artifact@v7 with: - node-version: "24" - cache: "npm" - cache-dependency-path: "tools/ui/package-lock.json" + name: llama-ui.zip + path: tools/ui/dist - name: ccache uses: ggml-org/ccache-action@v1.2.21 @@ -508,7 +476,6 @@ jobs: -DGGML_OPENVINO=ON \ -DCMAKE_INSTALL_RPATH='$ORIGIN' \ -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \ - -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \ ${{ env.CMAKE_ARGS }} cmake --build build/ReleaseOV --config Release --parallel @@ -552,7 +519,7 @@ jobs: key: release-ubuntu-24.04-openvino-release-no-preset-v1 windows-openvino: - needs: [check-release] + needs: [check-release, ui-build] if: ${{ needs.check-release.outputs.should_release == 'true' }} runs-on: windows-2022 @@ -562,8 +529,8 @@ jobs: env: # Sync versions in build-openvino.yml, build-self-hosted.yml, release.yml, build-cache.yml, .devops/openvino.Dockerfile - OPENVINO_VERSION_MAJOR: "2026.3" - OPENVINO_VERSION_FULL: "2026.3.0.22451.bd8d6542e3c" + OPENVINO_VERSION_MAJOR: "2026.3.1" + OPENVINO_VERSION_FULL: "2026.3.1.22476.56d9685302d" steps: - name: Set OpenVINO version output @@ -577,12 +544,11 @@ jobs: with: fetch-depth: 0 - - name: Setup Node.js - uses: actions/setup-node@v6 + - name: Download UI build + uses: actions/download-artifact@v7 with: - node-version: "24" - cache: "npm" - cache-dependency-path: "tools/ui/package-lock.json" + name: llama-ui.zip + path: tools/ui/dist - name: ccache uses: ggml-org/ccache-action@v1.2.21 @@ -682,7 +648,7 @@ jobs: windows-cpu: name: windows-cpu / ${{ matrix.arch }} - needs: [check-release] + needs: [check-release, ui-build] if: ${{ needs.check-release.outputs.should_release == 'true' }} runs-on: windows-2025-vs2026 @@ -702,12 +668,11 @@ jobs: with: fetch-depth: 0 - - name: Setup Node.js - uses: actions/setup-node@v6 + - name: Download UI build + uses: actions/download-artifact@v7 with: - node-version: "24" - cache: "npm" - cache-dependency-path: "tools/ui/package-lock.json" + name: llama-ui.zip + path: tools/ui/dist - name: Install Ninja run: | @@ -749,6 +714,8 @@ jobs: with: key: release-windows-2025-vs2026-${{ matrix.arch }}-cpu + # note: builds only the ggml-hip backend - llama-server is injected from the + # windows-cpu zip during the release "Merge artifacts" step windows-rocm: needs: [check-release] if: ${{ needs.check-release.outputs.should_release == 'true' }} @@ -769,6 +736,10 @@ jobs: with: fetch-depth: 0 + - name: Install Ninja + run: | + choco install ninja + - name: ccache uses: ggml-org/ccache-action@v1.2.21 with: @@ -822,33 +793,28 @@ jobs: - name: Build run: | - mkdir build - cd build - cmake .. ` - -G "Unix Makefiles" ` + cmake -S . -B build ` + -G "Ninja Multi-Config" ` -DCMAKE_PREFIX_PATH="${env:HIP_PATH}" ` - -DCMAKE_BUILD_TYPE=Release ` -DGGML_BACKEND_DL=ON ` -DGGML_NATIVE=OFF ` - -DGGML_CPU=ON ` - -DGGML_CPU_ALL_VARIANTS=ON ` + -DGGML_CPU=OFF ` -DGGML_HIP=ON ` -DCMAKE_C_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" ` -DCMAKE_CXX_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang++.exe" ` -DCMAKE_C_FLAGS="-Wno-error=incompatible-pointer-types" ` -DCMAKE_HIP_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" ` -DHIP_PATH="${env:HIP_PATH}" ` - -DGGML_HIP_ROCWMMA_FATTN=ON ` -DAMDGPU_TARGETS="${{ matrix.gpu_targets }}" - cmake --build . --config Release --parallel ${env:NUMBER_OF_PROCESSORS} + cmake --build build --config Release --parallel ${env:NUMBER_OF_PROCESSORS} --target ggml-hip - name: Verify HIP backend was built run: | - $hipDll = Get-ChildItem -Path build\bin -Filter "ggml-hip*.dll" -ErrorAction SilentlyContinue + $hipDll = Get-ChildItem -Path build\bin\Release -Filter "ggml-hip*.dll" -ErrorAction SilentlyContinue if (-not $hipDll) { Write-Host "##[error]ggml-hip*.dll was NOT produced. The HIP backend silently failed to build." - Write-Host "Contents of build\bin:" - Get-ChildItem build\bin | Format-Table -AutoSize + Write-Host "Contents of build\bin\Release:" + Get-ChildItem build\bin\Release | Format-Table -AutoSize exit 1 } Write-Host "HIP backend artifact found:" @@ -863,10 +829,40 @@ jobs: $rocmVersionShort = ('${{ matrix.ROCM_VERSION }}'.Split('.')[0..1] -join '.') echo "ROCM_VERSION_SHORT=$rocmVersionShort" >> $env:GITHUB_ENV + - name: Bundle HIP runtime DLLs (amdhip64_7.dll, rocm_kpack.dll, amd_comgr.dll) + run: | + $ErrorActionPreference = "Stop" + # See issue https://github.com/ggml-org/llama.cpp/issues/26929. + # ggml-hip.dll loads amdhip64_7.dll at run time. The Adrenalin driver + # ships an amdhip64_7.dll in System32, which the loader searches before PATH, + # so a matching DLL from PATH cannot win. Copy amdhip64 next to the + # binaries (exe directory is searched before System32) so the correct + # runtime is used. rocm_kpack.dll is amdhip64_7's direct dependency, so + # copy the matching version too. amd_comgr is copied as well to keep it + # in sync with the bundled amdhip64, avoiding a version mismatch with a + # amd_comgr from System32. + # rocblas/hipblaslt kernels resolve fine via PATH and are not copied. + $binPath = (rocm-sdk path --bin).Trim() + if (-not $binPath) { throw "rocm-sdk path --bin returned empty" } + write-host "ROCm bin path: $binPath" + + $patterns = @("amdhip64_7.dll", "rocm_kpack.dll", "amd_comgr.dll") + foreach ($pattern in $patterns) { + $files = Get-ChildItem -Path $binPath -Filter $pattern -ErrorAction SilentlyContinue + if (-not $files) { throw "no match for $pattern in $binPath" } + foreach ($f in $files) { + Copy-Item $f.FullName -Destination build\bin\Release -Force + write-host " copied $($f.Name)" + } + } + - name: Pack artifacts run: | - cp "LICENSE" "build\bin\" - 7z a -snl llama-bin-win-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.zip .\build\bin\* + 7z a -snl llama-bin-win-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.zip ` + .\build\bin\Release\ggml-hip.dll ` + .\build\bin\Release\amdhip64_7.dll ` + .\build\bin\Release\rocm_kpack.dll ` + .\build\bin\Release\amd_comgr.dll - name: Upload artifacts uses: actions/upload-artifact@v6 @@ -879,6 +875,8 @@ jobs: with: key: windows-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }} + # note: builds only the backend library - llama-server (with the embedded UI) + # is injected from the windows-cpu zip during the release "Merge artifacts" step windows: needs: [check-release] if: ${{ needs.check-release.outputs.should_release == 'true' }} @@ -909,13 +907,6 @@ jobs: id: checkout uses: actions/checkout@v6 - - name: Setup Node.js - uses: actions/setup-node@v6 - with: - node-version: "24" - cache: "npm" - cache-dependency-path: "tools/ui/package-lock.json" - - name: Install Vulkan SDK id: get_vulkan if: ${{ matrix.backend == 'vulkan' }} @@ -978,6 +969,8 @@ jobs: path: llama-bin-win-${{ matrix.backend }}-${{ matrix.arch }}.zip name: llama-bin-win-${{ matrix.backend }}-${{ matrix.arch }}.zip + # note: builds only the ggml-cuda backend - llama-server is injected from the + # windows-cpu zip during the release "Merge artifacts" step windows-cuda: name: windows-cuda (${{ matrix.cuda }}, ${{ matrix.arch }}) needs: [check-release] @@ -1006,13 +999,6 @@ jobs: id: checkout uses: actions/checkout@v6 - - name: Setup Node.js - uses: actions/setup-node@v6 - with: - node-version: "24" - cache: "npm" - cache-dependency-path: "tools/ui/package-lock.json" - - name: Install Cuda Toolkit uses: ./.github/actions/windows-setup-cuda with: @@ -1084,6 +1070,8 @@ jobs: with: key: release-windows-2022-${{ matrix.arch }}-cuda-${{ matrix.cuda }} + # note: builds only the ggml-sycl backend - llama-server is injected from the + # windows-cpu zip during the release "Merge artifacts" step windows-sycl: needs: [check-release] if: ${{ needs.check-release.outputs.should_release == 'true' }} @@ -1118,13 +1106,6 @@ jobs: Expand-Archive -Path "level-zero-win-sdk.zip" -DestinationPath "C:/level-zero-sdk" -Force "LEVEL_ZERO_V1_SDK_PATH=C:/level-zero-sdk" | Out-File -FilePath $env:GITHUB_ENV -Append - - name: Setup Node.js - uses: actions/setup-node@v6 - with: - node-version: "24" - cache: "npm" - cache-dependency-path: "tools/ui/package-lock.json" - - name: ccache uses: ggml-org/ccache-action@v1.2.21 with: @@ -1195,7 +1176,7 @@ jobs: key: release-windows-2022-x64-sycl ubuntu-24-sycl: - needs: [check-release] + needs: [check-release, ui-build] if: ${{ needs.check-release.outputs.should_release == 'true' }} strategy: @@ -1237,12 +1218,11 @@ jobs: wget -q "https://github.com/oneapi-src/level-zero/releases/download/v${LEVEL_ZERO_VERSION}/level-zero-devel_${LEVEL_ZERO_VERSION}%2B${LEVEL_ZERO_UBUNTU_VERSION}_amd64.deb" -O level-zero-devel.deb sudo apt-get install -y ./level-zero.deb ./level-zero-devel.deb - - name: Setup Node.js - uses: actions/setup-node@v6 + - name: Download UI build + uses: actions/download-artifact@v7 with: - node-version: "24" - cache: "npm" - cache-dependency-path: "tools/ui/package-lock.json" + name: llama-ui.zip + path: tools/ui/dist - name: ccache uses: ggml-org/ccache-action@v1.2.21 @@ -1288,7 +1268,7 @@ jobs: key: release-ubuntu-24.04-sycl-${{ matrix.build }} ubuntu-24-rocm: - needs: [check-release, get-version] + needs: [check-release, ui-build] if: ${{ needs.check-release.outputs.should_release == 'true' }} runs-on: ubuntu-24.04 @@ -1310,12 +1290,11 @@ jobs: with: fetch-depth: 0 - - name: Setup Node.js - uses: actions/setup-node@v6 + - name: Download UI build + uses: actions/download-artifact@v7 with: - node-version: "24" - cache: "npm" - cache-dependency-path: "tools/ui/package-lock.json" + name: llama-ui.zip + path: tools/ui/dist - name: Free up disk space uses: ggml-org/free-disk-space@v1.3.1 @@ -1388,7 +1367,6 @@ jobs: -DGPU_TARGETS="${{ matrix.gpu_targets }}" \ -DGGML_HIP=ON \ -DHIP_PLATFORM=amd \ - -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \ ${{ env.CMAKE_ARGS }} cmake --build build --config Release -j $(nproc) @@ -1417,7 +1395,7 @@ jobs: key: release-ubuntu-24.04-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }} ios-xcode: - needs: [check-release, get-version] + needs: [check-release] if: ${{ needs.check-release.outputs.should_release == 'true' }} runs-on: macos-26 @@ -1445,8 +1423,7 @@ jobs: -DLLAMA_BUILD_SERVER=OFF \ -DCMAKE_SYSTEM_NAME=iOS \ -DCMAKE_OSX_DEPLOYMENT_TARGET=16.0 \ - -DCMAKE_XCODE_ATTRIBUTE_DEVELOPMENT_TEAM=ggml \ - -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} + -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 @@ -1569,11 +1546,9 @@ jobs: # name: llama-bin-${{ matrix.chip_type }}-openEuler-${{ matrix.arch }}${{ matrix.use_acl_graph == 'on' && '-aclgraph' || '' }}.tar.gz ui-build: - needs: [check-release, get-version] + needs: [check-release] if: ${{ needs.check-release.outputs.should_release == 'true' }} uses: ./.github/workflows/ui-build.yml - with: - hf_ui_version: ${{ needs.get-version.outputs.ui_version }} release: if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }} @@ -1588,7 +1563,6 @@ jobs: runs-on: ubuntu-slim needs: - - get-version - windows - windows-cpu - windows-cuda @@ -1628,24 +1602,27 @@ jobs: path: ./artifact merge-multiple: true - - name: Move artifacts + - name: Merge artifacts id: move_artifacts run: | mkdir -p release - echo "Adding CPU backend files to existing zips..." + # the windows-cpu zip contains the full toolset (llama-server with the embedded + # UI, ggml-cpu) - inject it into the other windows zips so that every archive + # ships the same binaries, only with a different backend library on top + echo "Injecting windows-cpu binaries (llama-server + CPU backend) into the backend zips..." for arch in x64 arm64; do cpu_zip="artifact/llama-bin-win-cpu-${arch}.zip" temp_dir=$(mktemp -d) - echo "Extracting CPU backend for $arch..." + echo "Extracting windows-cpu-${arch} package..." unzip "$cpu_zip" -d "$temp_dir" - echo "Adding CPU files to $arch zips..." + echo "Merging into $arch zips..." for target_zip in artifact/llama-bin-win-*-${arch}.zip; do if [[ "$target_zip" == "$cpu_zip" ]]; then continue fi - echo "Adding CPU backend to $(basename "$target_zip")" + echo "Injecting into $(basename "$target_zip")" realpath_target_zip=$(realpath "$target_zip") (cd "$temp_dir" && zip -r "$realpath_target_zip" .) done @@ -1669,7 +1646,7 @@ jobs: id: download_ui uses: actions/download-artifact@v7 with: - name: ui-build + name: llama-ui.zip path: ./ui-dist - name: Package UI diff --git a/.github/workflows/server-sanitize.yml b/.github/workflows/server-sanitize.yml index 5d696282c70..a0ffc70c121 100644 --- a/.github/workflows/server-sanitize.yml +++ b/.github/workflows/server-sanitize.yml @@ -73,13 +73,6 @@ jobs: fetch-depth: 0 ref: ${{ github.event.inputs.sha || github.event.pull_request.head.sha || github.sha || github.head_ref || github.ref_name }} - - name: Setup Node.js - uses: actions/setup-node@v6 - with: - node-version: "24" - cache: "npm" - cache-dependency-path: "tools/ui/package-lock.json" - - name: Build id: cmake_build run: | diff --git a/.github/workflows/ui-build-self-hosted.yml b/.github/workflows/ui-build-self-hosted.yml index 7b7f8b60025..390a2f35f79 100644 --- a/.github/workflows/ui-build-self-hosted.yml +++ b/.github/workflows/ui-build-self-hosted.yml @@ -31,6 +31,6 @@ jobs: - name: Upload built UI uses: actions/upload-artifact@v6 with: - name: ui-build + name: llama-ui.zip path: tools/ui/dist/ retention-days: 1 diff --git a/.github/workflows/ui-build.yml b/.github/workflows/ui-build.yml index 85642f3f4b7..3fbd90c11cf 100644 --- a/.github/workflows/ui-build.yml +++ b/.github/workflows/ui-build.yml @@ -3,8 +3,8 @@ name: UI Build on: workflow_call: inputs: - hf_ui_version: - description: 'Version string for version.json (e.g. 12345)' + ui_version: + description: 'Version string embedded in build.json (e.g. b1234); defaults to b' required: false type: string @@ -17,6 +17,17 @@ jobs: steps: - name: Checkout code uses: actions/checkout@v6 + with: + fetch-depth: 0 + + - name: Resolve UI version + id: version + run: | + version="${{ inputs.ui_version }}" + if [ -z "$version" ]; then + version="b$(git rev-list --count HEAD)" + fi + echo "ui_version=${version}" >> $GITHUB_OUTPUT - name: Setup Node.js uses: actions/setup-node@v6 @@ -31,8 +42,7 @@ jobs: - name: Build application env: - HF_UI_VERSION: ${{ inputs.hf_ui_version || '' }} - LLAMA_BUILD_NUMBER: ${{ inputs.hf_ui_version || 'b0000' }} + LLAMA_BUILD_NUMBER: ${{ steps.version.outputs.ui_version }} run: npm run build working-directory: tools/ui @@ -43,6 +53,6 @@ jobs: - name: Upload built UI uses: actions/upload-artifact@v6 with: - name: ui-build + name: llama-ui.zip path: tools/ui/dist/ retention-days: 1 diff --git a/.github/workflows/ui-publish.yml b/.github/workflows/ui-publish.yml index 99a6d8420ff..e64ef32f801 100644 --- a/.github/workflows/ui-publish.yml +++ b/.github/workflows/ui-publish.yml @@ -37,7 +37,7 @@ jobs: - name: Download UI build artifact uses: actions/download-artifact@v7 with: - name: ui-build + name: llama-ui.zip path: tools/ui/dist/ - name: Create distribution archive diff --git a/.github/workflows/ui-self-hosted.yml b/.github/workflows/ui-self-hosted.yml index 79d7800d6bb..63521ead2d6 100644 --- a/.github/workflows/ui-self-hosted.yml +++ b/.github/workflows/ui-self-hosted.yml @@ -64,7 +64,7 @@ jobs: - name: Download built UI artifacts uses: actions/download-artifact@v6 with: - name: ui-build + name: llama-ui.zip path: tools/ui/dist/ - name: Run type checking @@ -106,7 +106,7 @@ jobs: - name: Download built UI artifacts uses: actions/download-artifact@v6 with: - name: ui-build + name: llama-ui.zip path: tools/ui/dist/ - name: Build Storybook diff --git a/.github/workflows/ui.yml b/.github/workflows/ui.yml index fa99a0cda2c..00a0804af38 100644 --- a/.github/workflows/ui.yml +++ b/.github/workflows/ui.yml @@ -63,7 +63,7 @@ jobs: - name: Download built UI artifacts uses: actions/download-artifact@v6 with: - name: ui-build + name: llama-ui.zip path: tools/ui/dist/ - name: Install dependencies @@ -126,7 +126,7 @@ jobs: - name: Download built UI artifacts (reuses ui-build) uses: actions/download-artifact@v6 with: - name: ui-build + name: llama-ui.zip path: tools/ui/dist/ - name: Install Playwright browsers diff --git a/CMakeLists.txt b/CMakeLists.txt index 730d5561fda..1d4bcf45f50 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -134,8 +134,8 @@ option(LLAMA_BUILD_TOOLS "llama: build tools" option(LLAMA_BUILD_EXAMPLES "llama: build examples" ${LLAMA_STANDALONE}) option(LLAMA_BUILD_SERVER "llama: build server example" ${LLAMA_STANDALONE}) option(LLAMA_BUILD_APP "llama: build the unified binary" ${LLAMA_STANDALONE}) -option(LLAMA_BUILD_UI "llama: build the embedded Web UI for server" ON) -option(LLAMA_USE_PREBUILT_UI "llama: use prebuilt UI from HF Bucket when available (requires LLAMA_BUILD_UI=ON)" ON) +option(LLAMA_BUILD_UI "llama: build the embedded Web UI for server" OFF) +option(LLAMA_USE_PREBUILT_UI "llama: use prebuilt UI from HF Bucket when available" ON) option(LLAMA_TOOLS_INSTALL "llama: install tools" ${LLAMA_TOOLS_INSTALL_DEFAULT}) option(LLAMA_TESTS_INSTALL "llama: install tests" ON) diff --git a/ci/run.sh b/ci/run.sh index 1f1e4bc033c..1701bc7ed05 100755 --- a/ci/run.sh +++ b/ci/run.sh @@ -189,8 +189,8 @@ if [ ! -z ${GG_BUILD_OPENVINO} ]; then fi CMAKE_EXTRA="${CMAKE_EXTRA} -DGGML_OPENVINO=ON" - # TODO: fix and re-enable the `test-llama-archs` test below - CTEST_EXTRA="-E test-llama-archs|test-recurrent-state-rollback-nemotron-h" + # TODO: fix and re-enable the `test-llama-archs` and `test-recurrent-state-rollback*` + CTEST_EXTRA="-E test-llama-archs|^test-recurrent-state-rollback" fi ## helpers diff --git a/common/arg.cpp b/common/arg.cpp index 86f8610a56d..4469612cd5b 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1643,6 +1643,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex } } ).set_env("LLAMA_ARG_CTX_SIZE")); + add_opt(common_arg( + { "--kv-unified-per-slot" }, "N", + "context limit per parallel slot (default: unset, behavior unchanged).\n" + "when set without -c/--ctx-size, the shared KV pool is sized to n_parallel*N", + [](common_params & params, int value) { + params.kv_unified_per_slot = value; + } + ).set_env("LLAMA_ARG_KV_UNIFIED_PER_SLOT").set_examples({ LLAMA_EXAMPLE_SERVER })); add_opt(common_arg( {"-n", "--predict", "--n-predict"}, "N", string_format( @@ -2644,6 +2652,27 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.mtmd_batch_max_tokens = value; } ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MTMD_BATCH_MAX_TOKENS")); + add_opt(common_arg( + {"--video-fps"}, "N", + string_format("target video frame rate (default: %.1f)", params.video_fps), + [](common_params & params, const std::string & value) { + params.video_fps = std::stof(value); + } + ).set_examples(mmproj_examples).set_env("LLAMA_ARG_VIDEO_FPS")); + add_opt(common_arg( + {"--video-timestamp-interval"}, "N", + string_format("interval in milliseconds between text timestamps (default: %" PRId64 ")", params.video_timestamp_interval_ms), + [](common_params & params, int value) { + params.video_timestamp_interval_ms = value; + } + ).set_examples(mmproj_examples).set_env("LLAMA_ARG_VIDEO_TIMESTAMP_INTERVAL")); + add_opt(common_arg( + {"--video-ffmpeg-dir"}, "DIR", + "path to the directory containing ffmpeg and ffprobe (default: search in PATH)", + [](common_params & params, const std::string & value) { + params.video_ffmpeg_bin_dir = value; + } + ).set_examples(mmproj_examples).set_env("LLAMA_ARG_VIDEO_FFMPEG_DIR")); if (params.is_gen_docs || llama_supports_rpc()) { add_opt(common_arg( {"--rpc"}, "SERVERS", @@ -2699,6 +2728,19 @@ common_params_context common_params_parser_init(common_params & params, llama_ex else { throw std::invalid_argument("invalid value"); } } ).set_env("LLAMA_ARG_LOAD_MODE")); + add_opt(common_arg( + {"--tensor-read-lazy"}, "MODE", + "on-demand reading of certain tensors, for example per-layer embeddings (default: auto)\n" + "- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)\n" + "- auto: on, but only for tensors larger than 4 GiB\n" + "- off: always keep them resident", + [](common_params & params, const std::string & value) { + /**/ if (value == "on") { params.lazy_mode = LLAMA_LAZY_MODE_ON; } + else if (value == "auto") { params.lazy_mode = LLAMA_LAZY_MODE_AUTO; } + else if (value == "off") { params.lazy_mode = LLAMA_LAZY_MODE_OFF; } + else { throw std::invalid_argument("invalid value"); } + } + ).set_env("LLAMA_ARG_TENSOR_READ_LAZY")); add_opt(common_arg( {"--numa"}, "TYPE", "attempt optimizations that help on some NUMA systems\n" @@ -2750,14 +2792,20 @@ common_params_context common_params_parser_init(common_params & params, llama_ex if (value < 0) { throw std::invalid_argument("invalid value"); } - for (int i = 0; i < value; ++i) { - // keep strings alive and avoid leaking memory by storing them in a static vector - static std::list buft_overrides; - buft_overrides.push_back(llm_ffn_exps_block_regex(i)); - params.tensor_buft_overrides.push_back({buft_overrides.back().c_str(), ggml_backend_cpu_buffer_type()}); - } + llm_add_n_cpu_ffn_overrides(value, LLM_FFN_EXPS_REGEX, params.tensor_buft_overrides); } ).set_env("LLAMA_ARG_N_CPU_MOE")); + add_opt(common_arg( + {"-ncffn", "--n-cpu-ffn"}, "N", + "keep the dense FFN weights of the first N layers in the CPU\n" + "(dense models; for MoE expert weights use --n-cpu-moe)", + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("invalid value"); + } + llm_add_n_cpu_ffn_overrides(value, LLM_FFN_DENSE_REGEX, params.tensor_buft_overrides); + } + ).set_env("LLAMA_ARG_N_CPU_FFN")); GGML_ASSERT(params.n_gpu_layers < 0); // string_format would need to be extended for a default >= 0 add_opt(common_arg( {"-ngl", "--gpu-layers", "--n-gpu-layers"}, "N", @@ -4084,11 +4132,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex if (value < 0) { throw std::invalid_argument("invalid value"); } - for (int i = 0; i < value; ++i) { - static std::list buft_overrides_draft; - buft_overrides_draft.push_back(llm_ffn_exps_block_regex(i)); - params.speculative.draft.tensor_buft_overrides.push_back({buft_overrides_draft.back().c_str(), ggml_backend_cpu_buffer_type()}); - } + llm_add_n_cpu_ffn_overrides(value, LLM_FFN_EXPS_REGEX, params.speculative.draft.tensor_buft_overrides); } ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE")); @@ -4109,6 +4153,38 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.speculative.draft.n_min = value; } ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_MIN")); + add_opt(common_arg( + {"--spec-synth-len"}, "L", + "target mean synthetic acceptance length, including the target token (benchmarking only)", + [](common_params & params, const std::string & value) { + const std::string text = string_strip(value); + size_t pos = 0; + const double length = std::stod(text, &pos); + if (pos != text.size() || length == -1.0) { + throw std::invalid_argument("invalid value"); + } + params.speculative.synth_len = length; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_SYNTH_LEN")); + add_opt(common_arg( + {"--spec-synth-rates"}, "P0,P1,...", + "comma-separated unconditional per-position synthetic acceptance probabilities (benchmarking only)", + [](common_params & params, const std::string & value) { + const auto values = string_split(value, ','); + std::vector rates; + rates.reserve(values.size()); + for (const auto & raw : values) { + const std::string text = string_strip(raw); + size_t pos = 0; + const double rate = std::stod(text, &pos); + if (pos != text.size()) { + throw std::invalid_argument("invalid value"); + } + rates.push_back(rate); + } + params.speculative.synth_rates = std::move(rates); + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_SYNTH_RATES")); add_opt(common_arg( {"--spec-draft-p-split", "--draft-p-split"}, "P", diff --git a/common/common.cpp b/common/common.cpp index 3d54bd6002d..d162a38800e 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1688,6 +1688,7 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.main_gpu = params.main_gpu; mparams.split_mode = params.split_mode; mparams.load_mode = params.load_mode; + mparams.lazy_mode = params.lazy_mode; mparams.tensor_split = params.tensor_split; mparams.check_tensors = params.check_tensors; mparams.use_extra_bufts = !params.no_extra_bufts; diff --git a/common/common.h b/common/common.h index de49dac9f63..4e9448bb106 100644 --- a/common/common.h +++ b/common/common.h @@ -8,6 +8,7 @@ #include "ggml.h" #include "llama.h" +#include #include #include #include @@ -369,6 +370,9 @@ struct common_params_speculative_ngram_cache { struct common_params_speculative { std::vector types = { COMMON_SPECULATIVE_TYPE_NONE }; + double synth_len = -1.0; + std::vector synth_rates; + // used by Simple, MTP, Eagle3, etc. - all methods that require some kind of draft model common_params_speculative_draft draft; @@ -383,6 +387,10 @@ struct common_params_speculative { return !draft.mparams.empty(); } + bool has_synth() const { + return synth_len != -1.0 || !synth_rates.empty(); + } + uint32_t need_n_rs_seq() const { bool needs_rs_seq = std::any_of(types.begin(), types.end(), [&](auto t) { return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK; @@ -475,6 +483,8 @@ struct common_params { enum llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER; // how to split the model across GPUs enum llama_load_mode load_mode = LLAMA_LOAD_MODE_AUTO; // how to load the model + enum llama_lazy_mode lazy_mode = LLAMA_LAZY_MODE_AUTO; // on-demand reading of tensors marked by the arch + common_cpu_params cpuparams; common_cpu_params cpuparams_batch; @@ -589,6 +599,11 @@ struct common_params { int image_max_tokens = -1; int mtmd_batch_max_tokens = 1024; + // for video input + float video_fps = 4.0f; + int64_t video_timestamp_interval_ms = 5000; + std::string video_ffmpeg_bin_dir = ""; + // finetune struct lr_opt lr; enum ggml_opt_optimizer_type optimizer = GGML_OPT_OPTIMIZER_TYPE_ADAMW; @@ -612,6 +627,7 @@ struct common_params { bool cache_prompt = true; // whether to enable prompt caching bool cache_idle_slots = true; // save and clear idle slots upon starting a new task int32_t n_ctx_checkpoints = 32; // max number of context checkpoints per slot + int32_t kv_unified_per_slot = 0; // max context per parallel slot; 0 = unset int32_t checkpoint_min_step = 8192; // minimum spacing between context checkpoints int32_t cache_ram_mib = 8192; // -1 = no limit, 0 - disable, 1 = 1 MiB, etc. @@ -1108,19 +1124,30 @@ const char * const LLM_KV_SPLIT_TENSORS_COUNT = "split.tensors.count"; } // -// MoE utils +// FFN offload utils // const char * const LLM_FFN_EXPS_REGEX = "\\.ffn_(up|down|gate|gate_up)_(ch|)exps"; -inline std::string llm_ffn_exps_block_regex(int idx) { - return string_format("blk\\.%d%s", idx, LLM_FFN_EXPS_REGEX); +const char * const LLM_FFN_DENSE_REGEX = "\\.ffn_(up|down|gate)\\."; + +inline std::string llm_ffn_block_regex(int idx, const char * ffn_regex) { + return string_format("blk\\.%d%s", idx, ffn_regex); } inline llama_model_tensor_buft_override llm_ffn_exps_cpu_override() { return { LLM_FFN_EXPS_REGEX, ggml_backend_cpu_buffer_type() }; } +inline void llm_add_n_cpu_ffn_overrides(int n, const char * ffn_regex, std::vector & overrides) { + // keep strings alive and avoid leaking memory by storing them in a static list + static std::list buft_override_strings; + for (int i = 0; i < n; ++i) { + buft_override_strings.push_back(llm_ffn_block_regex(i, ffn_regex)); + overrides.push_back({buft_override_strings.back().c_str(), ggml_backend_cpu_buffer_type()}); + } +} + // // training utils // diff --git a/common/speculative.cpp b/common/speculative.cpp index 4eef2212e75..d34d1c9c595 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -14,6 +14,7 @@ #include #include +#include #include #include #include @@ -138,6 +139,7 @@ struct common_speculative_impl { const common_speculative_type type; uint32_t n_seq; + int32_t n_max; // maximum draft length after implementation-specific limits size_t n_call_begin = 0; // number of times this implementation was called for refresh. size_t n_call_draft = 0; // number of times this implementation was called for generation. @@ -157,7 +159,7 @@ struct common_speculative_impl { int64_t t_draft_us = 0; // total time spent in generating drafts in this implementation in microseconds. int64_t t_accept_us = 0; // total time spent in accumulation of this implementation in microseconds. - common_speculative_impl(common_speculative_type type, uint32_t n_seq) : type(type), n_seq(n_seq) {} + common_speculative_impl(common_speculative_type type, uint32_t n_seq, int32_t n_max) : type(type), n_seq(n_seq), n_max(n_max) {} virtual ~common_speculative_impl() = default; @@ -182,7 +184,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { std::vector smpls; common_speculative_impl_draft_simple(const common_params_speculative & params, uint32_t n_seq) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE, n_seq) + : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE, n_seq, params.draft.n_max) , params(params.draft) { auto * ctx_dft = this->params.ctx_dft; @@ -452,7 +454,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { std::vector g_embd_buf; common_speculative_impl_draft_eagle3(const common_params_speculative & params, uint32_t n_seq) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, n_seq) + : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, n_seq, params.draft.n_max) , params(params.draft) { SPC_TRC("%s", "adding speculative implementation 'draft-eagle3'\n"); @@ -923,12 +925,19 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { int32_t block_size = 0; llama_token mask_token_id = 0; + bool is_dflash2 = false; + bool is_mrope = false; + int32_t selector_top_k = 0; + // draft-dspark: the draft carries a Markov head and uses an anchor-first block layout const bool is_dspark; // dspark speculators bool sample_from_anchor = true; + // block-internal attention + bool causal_attn = false; + const int32_t * target_layer_ids = nullptr; // model_dft's extract layer indices uint32_t target_layer_ids_n = 0; @@ -937,7 +946,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { 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) + : common_speculative_impl(type, n_seq, params.draft.n_max) , params(params.draft) , is_dspark(type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) { @@ -966,9 +975,25 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { if (llama_model_meta_val_str(model_dft, "dflash.sample_from_anchor", buf, sizeof(buf)) >= 0) { sample_from_anchor = std::strcmp(buf, "true") == 0; } + if (llama_model_meta_val_str(model_dft, "dflash.attention.causal", buf, sizeof(buf)) >= 0) { + causal_attn = std::strcmp(buf, "true") == 0; + } } + + selector_top_k = llama_model_dflash_selector_top_k(model_dft); + is_dflash2 = selector_top_k > 0; mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft)); + if (is_dspark && this->params.p_min > 0.0f) { + char buf[16] = {}; + const bool has_conf = + llama_model_meta_val_str(model_dft, "dflash.has_confidence_head", buf, sizeof(buf)) < 0 || + std::strcmp(buf, "true") == 0; + if (!has_conf) { + throw std::runtime_error("DSpark draft has no confidence head: please set --spec-draft-p-min 0"); + } + } + LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(type).c_str()); LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min); LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u, sample_from_anchor=%s\n", __func__, @@ -983,10 +1008,18 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { this->params.n_max = std::min(this->params.n_max, n_draft_max); this->params.n_min = std::min(this->params.n_min, n_draft_max); } + this->n_max = this->params.n_max; batch = llama_batch_init(llama_n_batch(ctx_dft), 0, n_seq); batch_inject = llama_batch_init(llama_n_batch(ctx_dft), n_embd_dec, n_seq); + // embd batches on an M-RoPE draft need 4 position rows per token + is_mrope = llama_model_rope_type(model_dft) == LLAMA_ROPE_TYPE_MROPE; + if (is_mrope) { + free(batch_inject.pos); + batch_inject.pos = (llama_pos *) malloc(sizeof(llama_pos) * 4 * llama_n_batch(ctx_dft)); + } + smpls.resize(n_seq); for (auto & s : smpls) { common_params_sampling sparams; @@ -998,7 +1031,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { // offload draft sampling to the backend backend_chains.assign(n_seq, nullptr); - if (this->params.backend_sampling) { + if (this->params.backend_sampling && !is_dflash2) { for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { llama_sampler * chain = llama_sampler_chain_init(llama_sampler_chain_default_params()); llama_sampler_chain_add(chain, llama_sampler_init_top_k(10)); @@ -1017,8 +1050,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true); } - llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true); - llama_set_causal_attn(ctx_dft, false); // DFlash needs non-causal attention + // DFlash2 reads its selector lattice from h_nextn and never consumes raw logits. + llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ !is_dflash2); + llama_set_causal_attn(ctx_dft, causal_attn); // DFlash needs non-causal attention unless the model says otherwise } ~common_speculative_impl_draft_dflash() override { @@ -1118,11 +1152,24 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { } // fuse extracted features through DFlash encoder + // M-RoPE drafts read 4 position rows per token from embd batches, so pass them explicitly + std::vector enc_pos; + if (is_mrope) { + enc_pos.resize((size_t) 4 * n_chunk); + for (int32_t i = 0; i < n_chunk; ++i) { + const llama_pos p = batch_in.pos[i_batch_beg[seq_id] + offset + i]; + enc_pos[0 * n_chunk + i] = p; + enc_pos[1 * n_chunk + i] = p; + enc_pos[2 * n_chunk + i] = p; + enc_pos[3 * n_chunk + i] = 0; + } + } + llama_batch enc_batch = { /*.n_tokens =*/ n_chunk, /*.token =*/ nullptr, /*.embd =*/ features_buf.data(), - /*.pos =*/ nullptr, + /*.pos =*/ is_mrope ? enc_pos.data() : nullptr, /*.n_seq_id =*/ nullptr, /*.seq_id =*/ nullptr, /*.logits =*/ nullptr, @@ -1143,7 +1190,13 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { std::memcpy(batch_inject.embd, inp_g, (size_t) n_chunk * n_embd_dec * sizeof(float)); for (int32_t i = 0; i < n_chunk; ++i) { - batch_inject.pos[i] = batch_in.pos[i_batch_beg[seq_id] + offset + i]; + const llama_pos p = batch_in.pos[i_batch_beg[seq_id] + offset + i]; + batch_inject.pos[i] = p; + if (is_mrope) { + batch_inject.pos[1 * n_chunk + i] = p; + batch_inject.pos[2 * n_chunk + i] = p; + batch_inject.pos[3 * n_chunk + i] = 0; + } batch_inject.n_seq_id[i] = 1; batch_inject.seq_id[i][0] = seq_id; batch_inject.logits[i] = false; @@ -1186,7 +1239,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { i_block_beg[seq_id] = batch.n_tokens; n_block [seq_id] = n_block_tokens; for (int32_t i = 0; i < n_block_tokens; ++i) { - common_batch_add(batch, i == 0 ? dp.id_last : mask_token_id, n + i, { seq_id }, true); + common_batch_add(batch, i == 0 ? dp.id_last : mask_token_id, n + i, { seq_id }, !is_dflash2); } } @@ -1214,6 +1267,36 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { auto & result = *dp.result; + if (is_dflash2) { + const float * lattice = llama_get_embeddings_nextn(ctx_dft); + GGML_ASSERT(lattice && "DFlash2 selector produced no lattice"); + + int32_t predecessor = 0; + for (int32_t i = 1; i < n_block_tokens; ++i) { + const float * row = lattice + (size_t) (beg + i) * n_embd_dec; + const float * scores = row + selector_top_k + (size_t) predecessor * selector_top_k; + + predecessor = (int32_t) std::distance(scores, + std::max_element(scores, scores + selector_top_k)); + if (params.p_min > 0.0f) { + // softmax(scores) at the argmax, i.e. 1 / sum(exp(s_k - s_max)) + float sum = 0.0f; + for (int32_t k = 0; k < selector_top_k; ++k) { + sum += std::exp(scores[k] - scores[predecessor]); + } + if (1.0f / sum < params.p_min) { + break; + } + } + result.push_back((llama_token) row[predecessor]); + } + + if (result.size() < (size_t) params.n_min) { + result.clear(); + } + continue; + } + if (is_dspark) { // DSpark: read from the first draft slot, truncate below the confidence threshold const float * conf = params.p_min > 0.0f ? llama_get_embeddings_nextn(ctx_dft) : nullptr; @@ -1315,7 +1398,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { std::vector> chain_h; common_speculative_impl_draft_mtp(const common_params_speculative & params, uint32_t n_seq) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_MTP, n_seq) + : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_MTP, n_seq, params.draft.n_max) , params(params.draft) { auto * ctx_tgt = this->params.ctx_tgt; @@ -1382,6 +1465,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { c.reserve((size_t) (this->params.n_max + 1) * n_embd); } } + this->n_max = this->params.n_max; pending_h.assign(n_seq, std::vector(n_embd, 0.0f)); @@ -1726,7 +1810,7 @@ struct common_speculative_impl_ngram_simple : public common_speculative_impl { common_speculative_impl_ngram_simple( const common_params_speculative & params, uint32_t n_seq, common_ngram_simple_config config) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, n_seq) + : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, n_seq, params.ngram_simple.size_m) , params(params.ngram_simple) , config(config) { @@ -1770,7 +1854,7 @@ struct common_speculative_impl_ngram_map_k : public common_speculative_impl { const common_ngram_map & config, uint32_t n_seq) : common_speculative_impl(config.key_only ? COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K - : COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, n_seq) + : COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, n_seq, config.size_value) { for (uint32_t i = 0; i < n_seq; i++) { this->config.push_back(config); @@ -1841,7 +1925,7 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl { common_speculative_impl_ngram_mod( const common_params_speculative & params, uint32_t n_seq) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_MOD, n_seq) + : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_MOD, n_seq, params.ngram_mod.n_max) , params(params.ngram_mod) , mod(params.ngram_mod.n_match, 4*1024*1024) , verbose(std::getenv("LLAMA_TRACE") != nullptr) { @@ -2017,7 +2101,7 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl { const std::string & path_dynamic, bool save_dynamic, bool save_static) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE, n_seq) + : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE, n_seq, n_draft) , params(params.ngram_cache) , n_draft(n_draft) , save_dynamic(save_dynamic) @@ -2138,6 +2222,8 @@ struct common_speculative { // which implementaion was used for a given seq_id std::vector impl_last; + + std::vector synth_probs; }; static common_ngram_map get_common_ngram_map( @@ -2316,6 +2402,101 @@ int32_t common_speculative_n_max(const common_params_speculative * spec) { return n_max; } +int32_t common_speculative_n_max(const common_speculative * spec) { + int32_t n_max = 0; + + if (spec == nullptr) { + return n_max; + } + + for (const auto & impl : spec->impls) { + n_max = std::max(n_max, std::max(0, impl->n_max)); + } + + return n_max; +} + +std::vector common_speculative_synth_rates_resolve(const common_params_speculative * spec, int32_t n_max) { + const bool has_length = spec->synth_len != -1.0; + const bool has_rates = !spec->synth_rates.empty(); + + if (!has_length && !has_rates) { + return {}; + } + if (has_length && has_rates) { + throw std::invalid_argument("synthetic acceptance length and rates are mutually exclusive"); + } + + if (n_max <= 0) { + throw std::invalid_argument("synthetic acceptance requires at least one speculative token"); + } + + if (has_rates) { + const auto & rates = spec->synth_rates; + if (rates.size() != (size_t) n_max) { + throw std::invalid_argument(string_format( + "synthetic acceptance rates must contain %d values, got %zu", n_max, rates.size())); + } + + for (size_t i = 0; i < rates.size(); ++i) { + if (!std::isfinite(rates[i]) || rates[i] < 0.0 || rates[i] > 1.0) { + throw std::invalid_argument("synthetic acceptance rates must be finite and within [0, 1]"); + } + if (i > 0 && rates[i] > rates[i - 1]) { + throw std::invalid_argument("synthetic acceptance rates must be monotonically non-increasing"); + } + } + + return rates; + } + + const double length = spec->synth_len; + const double length_max = (double) n_max + 1.0; + if (!std::isfinite(length) || length < 1.0 || length > length_max) { + throw std::invalid_argument(string_format( + "synthetic acceptance length must be finite and within [1, %.0f]", length_max)); + } + + double p = 0.0; + if (length == length_max) { + p = 1.0; + } else if (length > 1.0) { + double p_min = 0.0; + double p_max = 1.0; + for (int i = 0; i < 32; ++i) { + const double p_mid = 0.5 * (p_min + p_max); + double sum = 0.0; + double term = p_mid; + for (int32_t j = 0; j < n_max; ++j) { + sum += term; + term *= p_mid; + } + + if (sum < length - 1.0) { + p_min = p_mid; + } else { + p_max = p_mid; + } + } + p = 0.5 * (p_min + p_max); + } + + std::vector rates; + rates.reserve(n_max); + double rate = p; + for (int32_t i = 0; i < n_max; ++i) { + rates.push_back(rate); + rate *= p; + } + + return rates; +} + +const std::vector & common_speculative_get_synth_probs(const common_speculative * spec) { + GGML_ASSERT(spec); + return spec->synth_probs; +} + common_params common_base_params_to_speculative(const common_params & params) { const bool has_draft = params.speculative.has_dft(); @@ -2568,13 +2749,39 @@ common_speculative * common_speculative_init(common_params_speculative & params, return nullptr; } - auto * result = new common_speculative { - /* .dparams = */ common_speculative_draft_params_vec(n_seq), - /* .impls = */ std::move(impls), - /* .impl_last = */ std::vector(n_seq, nullptr) - }; + common_speculative_ptr result(new common_speculative { + /* .dparams = */ common_speculative_draft_params_vec(n_seq), + /* .impls = */ std::move(impls), + /* .impl_last = */ std::vector(n_seq, nullptr), + /* .synth_probs = */ {}, + }); - return result; + const int32_t n_max_configured = common_speculative_n_max(¶ms); + const int32_t n_max_effective = common_speculative_n_max(result.get()); + const auto rates = common_speculative_synth_rates_resolve(¶ms, n_max_effective); + + std::vector rates_str; + rates_str.reserve(rates.size()); + result->synth_probs.reserve(rates.size()); + double rate_prev = 1.0; + double acceptance_length = 1.0; + for (const double rate : rates) { + result->synth_probs.push_back(rate_prev > 0.0 ? rate / rate_prev : 0.0); + rates_str.push_back(string_format("%.6g", rate)); + rate_prev = rate; + acceptance_length += rate; + } + if (!result->synth_probs.empty()) { + SPC_WRN("%s", "synthetic speculative acceptance is enabled for benchmarking; generated output is not valid\n"); + if (n_max_effective != n_max_configured) { + SPC_WRN("synthetic acceptance draft limit was reduced from %d to %d by the initialized speculative implementations\n", + n_max_configured, n_max_effective); + } + SPC_INF("synthetic acceptance: n_max = %zu, mean length = %.6f, rates = [%s]\n", + rates.size(), acceptance_length, string_join(rates_str, ", ").c_str()); + } + + return result.release(); } void common_speculative_free(common_speculative * spec) { diff --git a/common/speculative.h b/common/speculative.h index 12ae31b7de5..22505891f7e 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -26,6 +26,15 @@ std::string common_speculative_type_to_str(enum common_speculative_type type); // return the max number of draft tokens based on the speculative parameters int32_t common_speculative_n_max(const common_params_speculative * spec); +// return the max number of draft tokens from the initialized implementations +int32_t common_speculative_n_max(const common_speculative * spec); + +// validate and resolve the unconditional synthetic acceptance rates +std::vector common_speculative_synth_rates_resolve(const common_params_speculative * spec, int32_t n_max); + +// return the conditional synthetic acceptance probabilities +const std::vector & common_speculative_get_synth_probs(const common_speculative * spec); + common_params common_base_params_to_speculative(const common_params & params); struct common_speculative_output_limits { diff --git a/conversion/__init__.py b/conversion/__init__.py index 8de97e95969..a5632fcc4bb 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -54,6 +54,7 @@ "DeepseekV3ForCausalLM": "deepseek", "DeepseekV32ForCausalLM": "deepseek", "DFlashDraftModel": "qwen", + "DFlash2DraftModel": "qwen", "Qwen3DSparkModel": "qwen", "DSparkDraftModel": "qwen", "DSparkSpeculator": "qwen", @@ -235,6 +236,8 @@ "Qwen3_5ForConditionalGeneration": "qwen", "Qwen3_5MoeForCausalLM": "qwen", "Qwen3_5MoeForConditionalGeneration": "qwen", + "Qwen4ExpForCausalLM": "qwen4exp", + "Qwen4ExpForConditionalGeneration": "qwen4exp", "RND1": "qwen", "RWForCausalLM": "falcon", "RWKV6Qwen2ForCausalLM": "rwkv", @@ -332,6 +335,7 @@ "Qwen3VLMoeForConditionalGeneration": "qwen3vl", "Qwen3_5ForConditionalGeneration": "qwen3vl", "Qwen3_5MoeForConditionalGeneration": "qwen3vl", + "Qwen4ExpForConditionalGeneration": "qwen4exp", "RADIOModel": "nemotron", "Sarashina2VisionForCausalLM": "sarashina2", "SmolVLMForConditionalGeneration": "smolvlm", diff --git a/conversion/base.py b/conversion/base.py index 56547ace009..daae28e92ad 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -1006,12 +1006,16 @@ def load(): else: raise ValueError(f"Unknown file type: {self.ftype.name}") + # 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 = gguf.quants.quantize(data, data_qtype) + data = quantize(data_qtype) except gguf.QuantError as e: logger.warning("%s, %s", e, "falling back to F16") data_qtype = gguf.GGMLQuantizationType.F16 - data = gguf.quants.quantize(data, data_qtype) + data = quantize(data_qtype) shape = gguf.quant_shape_from_byte_shape(data.shape, data_qtype) if data.dtype == np.uint8 else data.shape diff --git a/conversion/nemotron.py b/conversion/nemotron.py index e5d16718511..07fbc653149 100644 --- a/conversion/nemotron.py +++ b/conversion/nemotron.py @@ -202,6 +202,10 @@ class NemotronHModel(GraniteHybridModel): is_moe: bool = False supports_mtp_export = True + _SSM_LAYER_TYPES = {"mamba", "linear_attention"} + _ATTN_LAYER_TYPES = {"attention", "full_attention"} + _MLP_LAYER_TYPES = {"moe"} + 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 @@ -242,8 +246,8 @@ def __init__(self, *args, **kwargs): self._ssm_layers = [i for i, val in enumerate(pattern) if val == "M"] self._mlp_layers = [i for i, val in enumerate(pattern) if val == ("E" if self.is_moe else "-")] else: - self._ssm_layers = [i for i, val in enumerate(pattern) if val == "mamba"] - self._mlp_layers = [i for i, val in enumerate(pattern) if val == "moe"] + self._ssm_layers = [i for i, val in enumerate(pattern) if val in self._SSM_LAYER_TYPES] + self._mlp_layers = [i for i, val in enumerate(pattern) if val in self._MLP_LAYER_TYPES] # `--no-mtp` drops it entirely; `--mtp` exports only the MTP head self._mtp_bid: int | None = None @@ -272,7 +276,7 @@ def get_attn_layers(self): if isinstance(pattern, str): return [i for i, val in enumerate(pattern) if val == "*"] - return [i for i, val in enumerate(pattern) if val == "attention"] + return [i for i, val in enumerate(pattern) if val in self._ATTN_LAYER_TYPES] @classmethod def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: @@ -298,6 +302,10 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca ) if not keep: return None + # PEFT names adapter tensors using model.layers.*, while Nemotron-H checkpoints + # and the GGUF tensor map use backbone.layers.* + if name.startswith("model.layers.") and ".mixer." in name: + name = name.replace("model.layers.", "backbone.layers.", 1) return super().filter_tensors((name, gen)) def prepare_metadata(self, vocab_only: bool): diff --git a/conversion/qwen.py b/conversion/qwen.py index cdba8a63e9c..419611896fc 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -639,7 +639,7 @@ class Qwen3_5MoeTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase): model_arch = gguf.MODEL_ARCH.QWEN35MOE -@ModelBase.register("DFlashDraftModel") +@ModelBase.register("DFlashDraftModel", "DFlash2DraftModel") @ModelBase.example("z-lab/Qwen3.5-9B-DFlash") class DFlashModel(Qwen3Model): model_arch = gguf.MODEL_ARCH.DFLASH @@ -678,34 +678,98 @@ def set_vocab(self): def set_gguf_parameters(self): super().set_gguf_parameters() - block_size = self.hparams.get("block_size", 16) - self.gguf_writer.add_block_size(block_size) dflash_config = self.hparams.get("dflash_config", {}) + block_size = dflash_config.get("block_size", self.hparams.get("block_size", 16)) + self.gguf_writer.add_block_size(block_size) + + if "conv_kernel_size" in dflash_config: + self.gguf_writer.add_conv_kernel_size(int(dflash_config["conv_kernel_size"])) + self.gguf_writer.add_conv_group_size(int(dflash_config["conv_group_size"])) + self.gguf_writer.add_selector_rank(int(dflash_config["selector_rank"])) + self.gguf_writer.add_selector_top_k(int(dflash_config["selector_top_k"])) + + output_multiplier = dflash_config.get( + "output_multiplier", self.hparams.get("output_multiplier") + ) + if output_multiplier is not None: + self.gguf_writer.add_logit_scale(float(output_multiplier)) + softcap = dflash_config.get( + "final_logit_softcapping", self.hparams.get("final_logit_softcapping") + ) + if softcap is not None and float(softcap) > 0: + self.gguf_writer.add_final_logit_softcapping(float(softcap)) + embedding_scale = dflash_config.get( + "input_embedding_scale", self.hparams.get("input_embedding_scale") + ) + if embedding_scale is not None: + self.gguf_writer.add_embedding_scale(float(embedding_scale)) target_layer_ids = dflash_config.get("target_layer_ids", []) if target_layer_ids: extract_layer_ids = [i + 1 for i in target_layer_ids] self.gguf_writer.add_target_layers(extract_layer_ids) - use_sliding_window = self.hparams.get("use_sliding_window", False) - sliding_window = self.hparams.get("sliding_window") + use_sliding_window = self.hparams.get("use_sliding_window", False) or dflash_config.get("use_swa", False) + sliding_window = dflash_config.get("swa_window_size") or self.hparams.get("sliding_window") layer_types = self.hparams.get("layer_types") if use_sliding_window and sliding_window and layer_types: is_swa = [lt == "sliding_attention" for lt in layer_types] self.gguf_writer.add_sliding_window(sliding_window) self.gguf_writer.add_sliding_window_pattern(is_swa) + causal = self.hparams.get("is_causal") + if causal is None: + causal = dflash_config.get("causal") + if causal is not None: + self.gguf_writer.add_causal_attention(bool(causal)) + + # M-RoPE target: the draft ropes on the temporal dim only, so write + # degenerate sections [n_rot/2, 0, 0, 0] + if self._target_uses_mrope(): + head_dim = self.hparams.get("head_dim") or self.hparams["hidden_size"] // self.hparams["num_attention_heads"] + self.gguf_writer.add_rope_dimension_sections([head_dim // 2, 0, 0, 0]) + + def _target_uses_mrope(self) -> bool: + if self.target_model_dir is None: + return False + with open(self.target_model_dir / "config.json", "r", encoding="utf-8") as f: + cfg = json.load(f) + cfg = cfg.get("text_config", cfg) + rope = cfg.get("rope_parameters") or cfg.get("rope_scaling") or {} + return "mrope_section" in rope + @classmethod def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: name, gen = item if not name.startswith("model."): name = "model." + name + if "sink" in name and not name.endswith(".weight"): + name += ".weight" return super().filter_tensors((name, gen)) + _ROPE_PERMUTE_SUFFIXES = ( + "self_attn.q_proj.weight", + "self_attn.k_proj.weight", + "self_attn.q_norm.weight", + "self_attn.k_norm.weight", + ) + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: if name == "model.embed_tokens.weight" and not self.hparams.get("has_embed_tokens", True): return + # interleaved-rope checkpoints (rope_is_neox_style = false) -> NeoX layout: per head, even dims first then odd + if not self.hparams.get("rope_is_neox_style", True) and name.endswith(self._ROPE_PERMUTE_SUFFIXES): + head_dim = self.hparams["head_dim"] + shape = data_torch.shape + data_torch = data_torch.reshape(-1, head_dim // 2, 2, *shape[1:]).transpose(1, 2).reshape(shape) + + if name in ( + "model.candidate_selector.predecessor_codebook", + "model.candidate_selector.successor_codebook", + ): + name += ".weight" + yield from super().modify_tensors(data_torch, name, bid) @@ -759,6 +823,10 @@ def set_gguf_parameters(self): super().set_gguf_parameters() self.gguf_writer.add_sample_from_anchor(self._sample_from_anchor) + # confidence head is optional: vanilla-markov exports ship without it + has_conf = any("confidence_head.proj" in name for name in self.model_tensors) + self.gguf_writer.add_has_confidence_head(has_conf) + @classmethod def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: if item[0] == "t2d": # not used at runtime @@ -777,7 +845,7 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter self._d2t = data_torch return - if self._n_vocab_draft == self.hparams["vocab_size"] and name.endswith(("embed_tokens.weight", "lm_head.weight")): + if self._n_vocab_draft == self.hparams["vocab_size"] and name.endswith("lm_head.weight"): return # interleaved-rope checkpoints (rope_is_neox_style = false) -> NeoX layout: per head, even dims first then odd diff --git a/conversion/qwen4exp.py b/conversion/qwen4exp.py new file mode 100644 index 00000000000..168796d616b --- /dev/null +++ b/conversion/qwen4exp.py @@ -0,0 +1,195 @@ +from __future__ import annotations + +from typing import Iterable, cast + +import torch +from torch import Tensor + +import gguf +import numpy as np + +from .base import ModelBase +from .qwen import _LinearAttentionVReorderBase, _Qwen35MRopeMixin +from .qwen3vl import Qwen3VLVisionModel + + +@ModelBase.register("Qwen4ExpForConditionalGeneration", "Qwen4ExpForCausalLM") +@ModelBase.example("Qwen/Qwen3.8-Flash-Next") +class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase): + """Qwen3.8-Flash-Next. + + Shares the Qwen3.5 gated delta net and interleaved mrope, and adds three things: + hyper-connections in place of every layer norm, QSA sparse attention on the full + attention layers, and PLE n-gram hash embeddings on a single layer. + """ + + model_arch = gguf.MODEL_ARCH.QWEN4EXP + + # the MTP block is a separate draft head; vLLM drops it too + supports_mtp_export = False + no_mtp = True + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # only the shard names, so the table itself is never held + self._ple_shards: dict[int, str] = {} + self._ple_row_dim: int | None = None + + def _read_hash_constants(self, suffix: str) -> list[int]: + """Read an int64 PLE constant straight from the checkpoint. + + prepare_tensors() casts every non-float dtype to float32 before + modify_tensors() sees it (base.py), which would silently round these + 45-bit multipliers. Reading the lazy tensor here bypasses that. + """ + for name, gen in self.model_tensors.items(): + if name.endswith(suffix): + t = gen() + if t.dtype != torch.int64: + t = t.to(torch.int64) + return [int(x) for x in t.tolist()] + raise ValueError(f"PLE constant {suffix!r} missing from the checkpoint") + + def set_gguf_parameters(self): + super().set_gguf_parameters() + hp = self.hparams + + self.gguf_writer.add_hyper_connection_count(hp["hc_count"]) + self.gguf_writer.add_hyper_connection_low_rank(hp["hc_lowrank"]) + + n_layer = hp["num_hidden_layers"] + self.gguf_writer.add_indexer_head_count(hp["indexer_n_heads"]) + self.gguf_writer.add_indexer_key_length(hp["indexer_head_dim"]) + self.gguf_writer.add_indexer_top_k(hp["indexer_budget"]) + ratio = hp["indexer_compress_ratio"] + layer_types = hp["layer_types"] + self.gguf_writer.add_attention_compress_ratios( + [ratio if layer_types[i] == "full_attention" else 0 for i in range(n_layer)] + ) + + # ple_layer_ids is 1-based in the HF config; empty means no n-gram table, + # so emit no PLE keys rather than optional ones + ple_layers = [i - 1 for i in hp["ple_layer_ids"]] + if not ple_layers: + return + self.gguf_writer.add_ple_layers(ple_layers) + self.gguf_writer.add_ple_ngram_size(hp["ngram_size"]) + self.gguf_writer.add_ple_heads_per_ngram(hp["heads_per_ngram"]) + self.gguf_writer.add_ple_conv_kernel(hp["ple_conv_kernel_size"]) + self.gguf_writer.add_ple_eos_token_id(self._eos_token_id()) + # an image is decoded as an embeddings-only batch, so the graph has no placeholder + # ids to hash; carry the id and let it stand in for those positions + _img = self._image_token_id() + if _img is not None: + self.gguf_writer.add_ple_image_token_id(int(_img)) + if self._ple_row_dim is not None: + self.gguf_writer.add_embedding_length_per_layer_input(self._ple_row_dim) + + self.gguf_writer.add_ple_layer_multipliers( + self._read_hash_constants("ple_embedding.layer_multipliers")) + self.gguf_writer.add_ple_head_offsets( + self._read_hash_constants("ple_embedding.ngram_heads_offsets")) + self.gguf_writer.add_ple_head_vocab_sizes( + self._read_hash_constants("ple_embedding.ngram_heads_vocab_sizes")) + + def _image_token_id(self) -> int | None: + img = self.hparams.get("image_token_id") + return None if img is None else int(img) + + def _eos_token_id(self) -> int: + eos = self.hparams.get("eos_token_id") + if isinstance(eos, list): + # the PLE hash resets n-grams on the primary EOS + return int(eos[-1]) + if eos is None: + raise ValueError("eos_token_id is required: the PLE hash resets its n-grams on it") + return int(eos) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # int64 hash constants must stay exact; 1-D tensors force F32, so use KV + if name.endswith("ple_embedding.layer_multipliers"): + self._ple_multipliers = [int(x) for x in data_torch.tolist()] + return [] + if name.endswith("ple_embedding.ngram_heads_offsets"): + self._ple_head_offsets = [int(x) for x in data_torch.tolist()] + return [] + if name.endswith("ple_embedding.ngram_heads_vocab_sizes"): + self._ple_head_vocab_sizes = [int(x) for x in data_torch.tolist()] + return [] + + if ".ngram_embedding.shard_" in name: + return self._place_ple_shard(data_torch, name) + + # one projection feeds indexer q and k; split it, as minimax-m3 does + if ".indexer.index_qk_proj.weight" in name: + n_q = self.hparams["indexer_n_heads"] * self.hparams["indexer_head_dim"] + q = data_torch[:n_q] + k = data_torch[n_q:] + return [ + (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_Q_PROJ, bid, ".weight"), q), + (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_K_PROJ, bid, ".weight"), k), + ] + + # Gemma zero-centred gammas the inherited norm.weight rule misses + if name.endswith((".ple.norm_key.weight", ".ple.norm_query.weight", ".ple.norm_conv.weight", + ".indexer.q_layernorm.weight", ".indexer.k_layernorm.weight")): + return [(self.map_tensor_name(name), data_torch + 1)] + + if name.endswith(".ple.conv1d.weight"): + return [(self.map_tensor_name(name), data_torch.squeeze())] + + return super().modify_tensors(data_torch, name, bid) + + # the shards concatenate into a tensor of well over 100 GB + # use LazyChunkedTensor here, a single shard resident at a time + def _place_ple_shard(self, data_torch: Tensor, name: str) -> Iterable[tuple[str, Tensor]]: + + idx = int(name.rpartition(".shard_")[2].partition(".")[0]) + n_parts = self.hparams["split_ngram_parts"] + + self._ple_shards[idx] = name + self._ple_row_dim = int(data_torch.shape[-1]) + + if len(self._ple_shards) < n_parts: + return [] + + # the checkpoint may yield the shards in any order, the row order is by index + shards = [self._ple_shards[i] for i in sorted(self._ple_shards)] + rows = 0 + for shard in shards: + shape = self.model_tensors[shard]().shape + if int(shape[-1]) != self._ple_row_dim: + raise ValueError( + f"PLE shard {shard} has row dim {int(shape[-1])}, expected {self._ple_row_dim}") + rows += int(shape[0]) + + table = gguf.LazyChunkedTensor( + [self._load_ple_shard(shard) for shard in shards], + shape=(rows, self._ple_row_dim), + dtype=np.float32, + ) + gguf_name = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.PER_LAYER_TOKEN_EMBD] + return [(gguf_name + ".weight", cast(Tensor, table))] + + def _load_ple_shard(self, name: str): + def load() -> np.ndarray: + from .base import LazyTorchTensor + + # a fresh lazy tensor every call, or to_eager() memoizes every shard + eager = LazyTorchTensor.to_eager(self.model_tensors[name]()) + return eager.to(torch.float32).contiguous().numpy() + return load + + def prepare_tensors(self): + super().prepare_tensors() + n_parts = self.hparams.get("split_ngram_parts", 0) + if self._ple_shards and len(self._ple_shards) != n_parts: + raise ValueError( + f"got {len(self._ple_shards)} PLE embedding shards, expected {n_parts}" + ) + + +@ModelBase.register("Qwen4ExpForConditionalGeneration") +@ModelBase.example("Qwen/Qwen3.8-Flash-Next") +class Qwen4ExpVisionModel(Qwen3VLVisionModel): + """The vision tower is an unmodified Qwen3-VL ViT.""" diff --git a/docs/backend/OPENVINO.md b/docs/backend/OPENVINO.md index 3cdf631cebc..9b43807d36b 100644 --- a/docs/backend/OPENVINO.md +++ b/docs/backend/OPENVINO.md @@ -22,8 +22,8 @@ The OpenVINO backend is implemented in `ggml/src/ggml-openvino` and provides a t - [0. Prerequisites](#0-prerequisites) - [1. Install OpenVINO Runtime](#1-install-openvino-runtime) - [2. Build llama.cpp with OpenVINO Backend](#2-build-llamacpp-with-openvino-backend) - - [Automated Ubuntu Build Script](#automated-ubuntu-build-script) - - [Automated Windows Build Script](#automated-windows-build-script) + - [Ubuntu Build Script](#ubuntu-build-script) + - [Windows Build Script](#windows-build-script) - [3. Download Sample Model](#3-download-sample-model) - [4. Run Inference with OpenVINO Backend](#4-run-inference-with-openvino-backend) - [5. Docker Build](#5-docker-build) @@ -96,7 +96,7 @@ Although, the validated models below were tested with `llama-cli` using the `Q4_ - **SL** = Stateless (`GGML_OPENVINO_STATEFUL_EXECUTION=0`) - **SF** = Stateful (`GGML_OPENVINO_STATEFUL_EXECUTION=1`) - Note: The NPU operates in stateless mode only. -- **Validation system:** Intel® Core™ Ultra 5 238V (Lunar Lake) | 32 GB RAM | Ubuntu 24.04 | Intel OpenCL GPU Driver 26.18.38308.1 | Intel NPU Driver 1.33.0. +- **Validation system:** Intel® Core™ Ultra 5 238V (Lunar Lake) | 32 GB RAM | Ubuntu 24.04 | Intel OpenCL GPU Driver 26.31.39395.13-0 | Intel NPU Driver 1.35.0. - See [Known Limitations](#known-limitations) for context on observed failures. | Model | CPU (SL / SF) | GPU (SL / SF) | NPU (SL) | @@ -105,27 +105,32 @@ Although, the validated models below were tested with `llama-cli` using the `Q4_ | [bartowski/Llama-3.2-3B-Instruct-Q4_K_M](https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | | [bartowski/Meta-Llama-3.1-8B-Instruct-Q4_K_M](https://huggingface.co/bartowski/Meta-Llama-3.1-8B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | | | | | | -| [Qwen/qwen2.5-1.5b-instruct-q4_k_m](https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | -| [Qwen/qwen2.5-coder-7b-instruct-q4_k_m](https://huggingface.co/Qwen/Qwen2.5-Coder-7B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | -| [bartowski/Qwen_Qwen3-0.6B-Q4_K_M](https://huggingface.co/bartowski/Qwen_Qwen3-0.6B-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | -| [bartowski/Qwen_Qwen3-1.7B-Q4_K_M](https://huggingface.co/bartowski/Qwen_Qwen3-1.7B-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | -| [Qwen/Qwen3-4B-Q4_K_M](https://huggingface.co/Qwen/Qwen3-4B-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | -| [lm-kit/Qwen3-8B-Q4_K_M](https://huggingface.co/lm-kit/qwen-3-8b-instruct-gguf) | ✓ / ✓ | ✓ / ✗ | ✓ | +| [Qwen/qwen2.5-1.5b-instruct-q4_k_m](https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [Qwen/qwen2.5-coder-7b-instruct-q4_k_m](https://huggingface.co/Qwen/Qwen2.5-Coder-7B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [bartowski/Qwen_Qwen3-0.6B-Q4_K_M](https://huggingface.co/bartowski/Qwen_Qwen3-0.6B-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [bartowski/Qwen_Qwen3-1.7B-Q4_K_M](https://huggingface.co/bartowski/Qwen_Qwen3-1.7B-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [Qwen/Qwen3-4B-Q4_K_M](https://huggingface.co/Qwen/Qwen3-4B-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [lm-kit/Qwen3-8B-Q4_K_M](https://huggingface.co/lm-kit/qwen-3-8b-instruct-gguf) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [bartowski/Qwen_Qwen3.5-0.8B-Q4_K_M](https://huggingface.co/bartowski/Qwen_Qwen3.5-0.8B-GGUF) | ✓ / ✗ | ✓ / ✗ | ✗ | +| [bartowski/Qwen_Qwen3.5-2B-Q4_K_M](https://huggingface.co/bartowski/Qwen_Qwen3.5-2B-GGUF) | ✓ / ✗ | ✓ / ✗ | ✗ | +| [bartowski/Qwen_Qwen3.5-4B-Q4_K_M](https://huggingface.co/bartowski/Qwen_Qwen3.5-4B-GGUF) | ✓ / ✗ | ✓ / ✗ | ✗ | +| [lmstudio-community/Qwen3.5-9B-Q4_K_M](https://huggingface.co/lmstudio-community/Qwen3.5-9B-GGUF) | ✓ / ✗ | ✓ / ✗ | ✗ | | | | | | -| [unsloth/gemma-3-4b-it-Q4_K_M](https://huggingface.co/unsloth/gemma-3-4b-it-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | -| [bartowski/google_gemma-4-E2B-it-Q4_K_M](https://huggingface.co/bartowski/google_gemma-4-E2B-it-GGUF) | ✓ / ✗ | ✓ / ✗ | ✓ | +| [unsloth/gemma-3-4b-it-Q4_K_M](https://huggingface.co/unsloth/gemma-3-4b-it-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [bartowski/google_gemma-4-E2B-it-Q4_K_M](https://huggingface.co/bartowski/google_gemma-4-E2B-it-GGUF) | ✓ / ✗ | ✓ / ✗ | ✗ | | [bartowski/google_gemma-4-E4B-it-Q4_K_M](https://huggingface.co/bartowski/google_gemma-4-E4B-it-GGUF) | ✓ / ✗ | ✓ / ✗ | ✓ | -| [bartowski/gemma-4-12B-it-Q4_K_M](https://huggingface.co/bartowski/gemma-4-12B-it-GGUF) | ✓ / ✗ | ✓ / ✗ | ✗ | +| [bartowski/gemma-4-12B-it-Q4_K_M](https://huggingface.co/bartowski/gemma-4-12B-it-GGUF) | ✓ / ✗ | ✓ / ✗ | ✓ | | | | | | -| [bartowski/Phi-3-mini-4k-instruct-Q4_K_M](https://huggingface.co/bartowski/Phi-3-mini-4k-instruct-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | -| [bartowski/Phi-3.5-mini-instruct-Q4_K_M](https://huggingface.co/bartowski/Phi-3.5-mini-instruct-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | +| [bartowski/Phi-3-mini-4k-instruct-Q4_K_M](https://huggingface.co/bartowski/Phi-3-mini-4k-instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [bartowski/Phi-3.5-mini-instruct-Q4_K_M](https://huggingface.co/bartowski/Phi-3.5-mini-instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [bartowski/microsoft_Phi-4-mini-instruct-Q4_K_M](https://huggingface.co/bartowski/microsoft_Phi-4-mini-instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | | | | | | | [bartowski/Mistral-7B-Instruct-v0.3-Q4_K_M](https://huggingface.co/bartowski/Mistral-7B-Instruct-v0.3-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | | [QuantFactory/Ministral-3b-instruct.Q4_K_M](https://huggingface.co/QuantFactory/Ministral-3b-instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | | [bartowski/Ministral-8B-Instruct-2410-Q4_K_M](https://huggingface.co/bartowski/Ministral-8B-Instruct-2410-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | | | | | | | [bartowski/DeepSeek-R1-Distill-Llama-8B-Q4_K_M](https://huggingface.co/bartowski/DeepSeek-R1-Distill-Llama-8B-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | -| [bartowski/DeepSeek-R1-Distill-Qwen-7B-Q4_K_M](https://huggingface.co/bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | +| [bartowski/DeepSeek-R1-Distill-Qwen-7B-Q4_K_M](https://huggingface.co/bartowski/DeepSeek-R1-Distill-Qwen-7B-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | | | | | | | [ibm-granite/granite-4.0-350m-Q4_K_M](https://huggingface.co/ibm-granite/granite-4.0-350m-GGUF) | ✓ / ✓ | ✗ / ✗ | ✓ | | [ibm-granite/granite-4.0-micro-Q4_K_M](https://huggingface.co/ibm-granite/granite-4.0-micro-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | @@ -133,10 +138,10 @@ Although, the validated models below were tested with `llama-cli` using the `Q4_ | [ibm-research/granite-3.2-8b-instruct-Q4_K_M](https://huggingface.co/ibm-research/granite-3.2-8b-instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | | | | | | | [HuggingFaceTB/smollm2-1.7b-instruct-q4_k_m](https://huggingface.co/HuggingFaceTB/SmolLM2-1.7B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | -| [openbmb/MiniCPM-V-2_6-Q4_K_M](https://huggingface.co/openbmb/MiniCPM-V-2_6-gguf) | ✓ / ✓ | ✓ / ✗ | ✓ | -| [bartowski/tencent_Hunyuan-7B-Instruct-Q4_K_M](https://huggingface.co/bartowski/tencent_Hunyuan-7B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | -| [LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct-Q4_K_M](https://huggingface.co/LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | -| [bartowski/prism-ml_Bonsai-8B-unpacked-Q4_K_M](https://huggingface.co/bartowski/prism-ml_Bonsai-8B-unpacked-GGUF) | ✓ / ✓ | ✓ / ✗ | ✓ | +| [openbmb/MiniCPM-V-2_6-Q4_K_M](https://huggingface.co/openbmb/MiniCPM-V-2_6-gguf) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [bartowski/tencent_Hunyuan-7B-Instruct-Q4_K_M](https://huggingface.co/bartowski/tencent_Hunyuan-7B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct-Q4_K_M](https://huggingface.co/LGAI-EXAONE/EXAONE-3.5-7.8B-Instruct-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | +| [bartowski/prism-ml_Bonsai-8B-unpacked-Q4_K_M](https://huggingface.co/bartowski/prism-ml_Bonsai-8B-unpacked-GGUF) | ✓ / ✓ | ✓ / ✓ | ✓ | | | | | | | [gpustack/bge-m3-Q4_K_M.gguf](https://huggingface.co/gpustack/bge-m3-GGUF) | ✓ | ✗ | ✗ | @@ -217,18 +222,18 @@ cmake --build build\ReleaseOV --parallel > [!NOTE] > The Windows install path is `C:\Intel\openvino` (no spaces) to avoid quoting problems some CMake/Ninja toolchains have with `C:\Program Files (x86)\...`. Adjust to wherever you installed OpenVINO Runtime. From `cmd`, run `C:\Intel\openvino\setupvars.bat`; from PowerShell, run `& "C:\Intel\openvino\setupvars.ps1"` instead. Once the build is finished you can launch the binaries from any `cmd` or `PowerShell` window after sourcing the matching `setupvars` script for that shell. -#### Automated Ubuntu Build Script +#### Ubuntu Build Script For Ubuntu24 users, the following shell script automates the prerequisite installs (build tools, OpenCL ICD), the OpenVINO Runtime download/extract/setup, and the Ninja-based llama.cpp build. -Save the following as `ubuntu-llamacpp-ov-install.sh` next to where you want the `llama.cpp` folder to land, then run it: +Save the following as `build-llamacpp-ov.sh` next to where you want the `llama.cpp` folder to land, then run it: ```bash -chmod +x ubuntu-llamacpp-ov-install.sh -./ubuntu-llamacpp-ov-install.sh +chmod +x build-llamacpp-ov.sh +./build-llamacpp-ov.sh ```
-Click to expand ubuntu-llamacpp-ov-install.sh +Click to expand build-llamacpp-ov.sh ```bash #!/usr/bin/env bash @@ -237,8 +242,8 @@ chmod +x ubuntu-llamacpp-ov-install.sh # ============================================ set -euo pipefail -OPENVINO_VERSION_MAJOR="2026.3" -OPENVINO_VERSION_FULL="2026.3.0.22451.bd8d6542e3c" +OPENVINO_VERSION_MAJOR="2026.3.1" +OPENVINO_VERSION_FULL="2026.3.1.22476.56d9685302d" SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" OPENVINO_INSTALL_DIR="/opt/intel/openvino_${OPENVINO_VERSION_MAJOR}" @@ -313,8 +318,9 @@ fi echo "============================================" echo "Configuring with CMake..." echo "============================================" -# shellcheck disable=SC1091 +set +u source "${OPENVINO_ROOT}/setupvars.sh" +set -u cmake -B build/ReleaseOV -G Ninja \ -DCMAKE_BUILD_TYPE=Release \ @@ -334,27 +340,27 @@ echo " ./build/ReleaseOV/bin/llama-cli -m model.gguf" ``` > [!NOTE] -> The script pins OpenVINO `2026.3` via the `OPENVINO_VERSION_MAJOR` / `OPENVINO_VERSION_FULL` variables at the top — edit them to track a different release. +> The script pins OpenVINO `2026.3.1` via the `OPENVINO_VERSION_MAJOR` / `OPENVINO_VERSION_FULL` variables at the top — edit them to track a different release.
-#### Automated Windows Build Script +#### Windows Build Script For Windows users, the following `.bat` script automates the prerequisite installs (Git, Ninja, CMake, Visual Studio 2022 Build Tools, vcpkg + OpenCL), the OpenVINO Runtime download/extract, and the Ninja-based llama.cpp build. -Save the following as `windows-llamacpp-ov-install.bat` next to where you want the `llama.cpp` to land, then run it from either **Command Prompt** or **PowerShell**: +Save the following as `build-llamacpp-ov.bat` next to where you want the `llama.cpp` to land, then run it from either **Command Prompt** or **PowerShell**: ```cmd :: Command Prompt -windows-llamacpp-ov-install.bat +build-llamacpp-ov.bat ``` ```powershell # PowerShell -.\windows-llamacpp-ov-install.bat +.\build-llamacpp-ov.bat ```
-Click to expand windows-llamacpp-ov-install.bat +Click to expand build-llamacpp-ov.bat ```bat @echo off @@ -364,8 +370,8 @@ REM ============================================ REM llama.cpp OpenVINO Build Script (Ninja) REM ============================================ -set "OPENVINO_VERSION_MAJOR=2026.3" -set "OPENVINO_VERSION_FULL=2026.3.0.22451.bd8d6542e3c" +set "OPENVINO_VERSION_MAJOR=2026.3.1" +set "OPENVINO_VERSION_FULL=2026.3.1.22476.56d9685302d" set "SCRIPT_DIR=%~dp0" set "VCPKG_DIR=C:\vcpkg" @@ -453,9 +459,6 @@ if exist "%OPENVINO_INSTALL_DIR%\setupvars.bat" ( ) REM Move the single top-level folder contents into the versioned install dir. - REM NOTE: delayed expansion (!VAR!) is required because the surrounding else( ... ) - REM block is parsed once up-front, so %OPENVINO_EXTRACTED% would expand to "" here - REM and xcopy would then treat "\*" as C:\* and fail with "Cannot perform a cyclic copy". set "OPENVINO_EXTRACTED=" for /d %%i in ("%OPENVINO_EXTRACT_TMP%\*") do set "OPENVINO_EXTRACTED=%%i" if not defined OPENVINO_EXTRACTED ( @@ -547,7 +550,7 @@ endlocal ``` > [!NOTE] -> The script pins OpenVINO `2026.3` via the `OPENVINO_VERSION_MAJOR` / `OPENVINO_VERSION_FULL` variables at the top — edit them to track a different release. From any new shell, source the matching `setupvars` script via the junction — `call "C:\Intel\openvino\setupvars.bat"` from `cmd`, or `& "C:\Intel\openvino\setupvars.ps1"` from PowerShell. If `winget` cannot register Visual Studio Build Tools on first run, install them once manually and re-run the script from an elevated **Developer Command Prompt for VS 2022**. +> The script pins OpenVINO `2026.3.1` via the `OPENVINO_VERSION_MAJOR` / `OPENVINO_VERSION_FULL` variables at the top — edit them to track a different release. From any new shell, source the matching `setupvars` script via the junction — `call "C:\Intel\openvino\setupvars.bat"` from `cmd`, or `& "C:\Intel\openvino\setupvars.ps1"` from PowerShell. If `winget` cannot register Visual Studio Build Tools on first run, install them once manually and re-run the script from an elevated **Developer Command Prompt for VS 2022**.
@@ -712,6 +715,7 @@ Boolean flags follow a uniform convention: set to a **positive integer** (e.g. ` | `GGML_OPENVINO_CACHE_DIR` | String | `not set` | Directory for OpenVINO model caching (recommended: `/tmp/ov_cache`). Enables model caching when set. **Not supported on NPU devices.** | | `GGML_OPENVINO_COMPILED_MODEL_CACHE_DIR` | String | `not set` | Directory for the frontend compiled-model cache. When set, OpenVINO compiled models are exported as blobs and imported on later runs to skip weight requantization, graph conversion, and compilation for matching single-graph models. | | `GGML_OPENVINO_PREFILL_CHUNK_SIZE`| Integer | `256` | Token chunk size for **NPU** prefill (NPU-only; ignored on CPU/GPU). Must be a positive integer; otherwise the default is used. | +| `GGML_OPENVINO_NPU_COMPILE_CONFIG` | String | `not set` | NPU-only compiler mode parameters forwarded to OpenVINO as `NPU_COMPILATION_MODE_PARAMS`, for example `optimization-level=3`. | | `GGML_OPENVINO_STATEFUL_EXECUTION`| Boolean | `0` | Enable stateful KV cache for better performance. Recommended on CPU, GPU. | | `GGML_OPENVINO_DISABLE_CACHE` | Boolean | `0` | Disable the in-process compiled-model / decoder cache (cache is on by default). Set to `1` to disable. | | `GGML_OPENVINO_DISABLE_KV_SLICE` | Boolean | `0` | Disable the KV-cache input-tensor slicing optimization (slicing is on by default on CPU/GPU). Set to `1` to disable. | @@ -725,9 +729,11 @@ Boolean flags follow a uniform convention: set to a **positive integer** (e.g. ` | `GGML_OPENVINO_DEBUG_INPUT` | Boolean | `0` | Enable input debugging and print input tensor info. | | `GGML_OPENVINO_DEBUG_OUTPUT` | Boolean | `0` | Enable output debugging and print output tensor info. | | `GGML_OPENVINO_PRINT_CGRAPH_TENSOR_ADDRESS` | Boolean | `0` | Print tensor address map once. | +| `GGML_OPENVINO_LOG_UNSUPPORTED_OPS`| Boolean | `0` | Log warning messages with tensor details and rejection reasons for any ops not supported by the OpenVINO backend. Emits at `WARN` level (requires `--log-verbosity >= 2`, enabled by default). | > [!NOTE] ->`GGML_OPENVINO_STATEFUL_EXECUTION` is an **Experimental** feature to allow stateful execution for managing the KV cache internally inside the OpenVINO model, improving performance on CPUs and GPUs. Stateful execution is not effective on NPUs, and not all models currently support this feature. This feature is experimental and has been validated only with the llama-simple, llama-cli, llama-bench, and llama-run applications and is recommended to enable for the best performance. Other applications, such as llama-server and llama-perplexity, are not yet supported. +> - `GGML_OPENVINO_STATEFUL_EXECUTION` is an **Experimental** feature to allow stateful execution for managing the KV cache internally inside the OpenVINO model, improving performance on CPUs and GPUs. Stateful execution is not effective on NPUs, and not all models currently support this feature. This feature is experimental and has been validated only with the llama-simple, llama-cli, llama-bench, and llama-run applications and is recommended to enable for the best performance. Other applications, such as llama-server and llama-perplexity, are not yet supported. +> - `GGML_OPENVINO_LOG_UNSUPPORTED_OPS` emits logs at `WARN` level (`GGML_LOG_WARN`), which requires application log verbosity `--log-verbosity >= 2` (or `-lv 2`). ### Example Usage diff --git a/docs/backend/snapdragon/CMakeUserPresets.json b/docs/backend/snapdragon/CMakeUserPresets.json index 848d735f1c5..afc73923237 100644 --- a/docs/backend/snapdragon/CMakeUserPresets.json +++ b/docs/backend/snapdragon/CMakeUserPresets.json @@ -8,7 +8,7 @@ "toolset": { "value": "host=x86_64", "strategy": "external" }, "cacheVariables": { "ANDROID_ABI": "arm64-v8a", - "ANDROID_PLATFORM": "android-31", + "ANDROID_PLATFORM": "android-34", "CMAKE_TOOLCHAIN_FILE": "$env{ANDROID_NDK_ROOT}/build/cmake/android.toolchain.cmake", "CMAKE_C_FLAGS": "-march=armv8.7a+fp16+dotprod+i8mm -fvectorize -ffp-model=fast -fno-finite-math-only -flto -D_GNU_SOURCE", "CMAKE_CXX_FLAGS": "-march=armv8.7a+fp16+dotprod+i8mm -fvectorize -ffp-model=fast -fno-finite-math-only -flto -D_GNU_SOURCE", diff --git a/docs/backend/snapdragon/README.md b/docs/backend/snapdragon/README.md index e9f0e215858..391c8bf230f 100644 --- a/docs/backend/snapdragon/README.md +++ b/docs/backend/snapdragon/README.md @@ -2,39 +2,47 @@ ## Setup -### Android +The cross-compilation toolchain images are provided by the +[Qualcomm Snapdragon Toolchain registry](https://github.com/snapdragon-toolchain). +These Docker images include the Android NDK, OpenCL SDK, Hexagon SDK, CMake, and the necessary cross-compilers: -The easiest way to build llama.cpp for a Snapdragon-based Android device is using the toolchain Docker image (see github.com/snapdragon-toolchain). -This image includes Android NDK, OpenCL SDK, Hexagon SDK, CMake, etc. +* **Android toolchain**: `ghcr.io/snapdragon-toolchain/arm64-android:v0.7` +* **Linux toolchain**: `ghcr.io/snapdragon-toolchain/arm64-linux:v0.7` -This method works on Linux, macOS, and Windows. macOS and Windows users should install Docker Desktop. +The unified build utility (`scripts/snapdragon/build.py`) automatically pulls +and orchestrates these containers to perform target compilation. +You only need to ensure that Docker (or Docker Desktop on macOS/Windows) is running on your host machine. +Specific setup, build, and installation details for Linux and Windows on Snapdragon platforms are documented in: +* [Linux on Snapdragon guide](linux.md) +* [Windows on Snapdragon guide](windows.md) -``` -~/src/llama.cpp$ docker run -it -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-android:v0.7 -[d]/> cd /workspace -``` +## How to Build + +### Using build.py script (Recommended) -Note: The rest of the **Android** build process assumes that you're running inside the toolchain container. +The easiest way to build llama.cpp is by using the `scripts/snapdragon/build.py` script. It automatically copies the CMake presets, +launches the correct compilation Docker container, builds the libraries and tools, +installs them, and optionally pushes them to your ADB device. -### Windows On Snapdragon +Build and deploy for Android target (accepts `android` or `adb` alias): +``` +$ ./scripts/snapdragon/build.py --target adb --push +``` -Native Windows 11 arm64 builds has the following tools dependencies: -- MS Visual Studio 2026 (Community Edition or Pro) - - MSVC arm64 standard and runtime libraries - - UCRT and Driver Kit -- LLVM core libraries and Clang compiler (winget) -- CMake, Git, Python (winget) -- Hexagon SDK Community Edition 6.6 or later (see windows.md) -- OpenCL SDK 2.3 or later (see windows.md) +Build and deploy for Linux target (accepts `linux` or `lnx` alias): +``` +$ ./scripts/snapdragon/build.py --target linux:user@host --push +``` -Note: The rest of the **Windows** build process assumes that you're running natively in Powershell. -Adapt below build commands accordingly. +### Manual CMake Build -## How to Build +Alternatively, you can build llama.cpp manually by entering the cross-compilation Docker container and running the CMake commands: -Let's build llama.cpp with CPU, OpenCL, and Hexagon backends via CMake presets: +```bash +# Start the cross-compilation container manually: +~/src/llama.cpp$ docker run -it --rm -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-android:v0.7 -``` +# Inside the container, build the project using presets: [d]/workspace> cp docs/backend/snapdragon/CMakeUserPresets.json . [d]/workspace> cmake --preset arm64-android-snapdragon-release -B build-snapdragon @@ -68,19 +76,19 @@ Preset CMake variables: To generate an installable "package" simply use cmake --install: ``` -[d]/workspace> cmake --install build-snapdragon --prefix pkg-snapdragon/llama.cpp +[d]/workspace> cmake --install build-snapdragon --prefix pkg-android/llama.cpp -- Install configuration: "Release" --- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-cpu.so --- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-opencl.so --- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-hexagon.so --- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-htp-v73.so --- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-htp-v75.so --- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-htp-v79.so --- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-htp-v81.so --- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml.so +-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-cpu.so +-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-opencl.so +-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-hexagon.so +-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-htp-v73.so +-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-htp-v75.so +-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-htp-v79.so +-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-htp-v81.so +-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml.so ... --- Installing: /workspace/pkg-snapdragon/llama.cpp/bin/llama-bench --- Installing: /workspace/pkg-snapdragon/llama.cpp/bin/llama-cli +-- Installing: /workspace/pkg-android/llama.cpp/bin/llama-bench +-- Installing: /workspace/pkg-android/llama.cpp/bin/llama-cli ... ``` @@ -91,14 +99,14 @@ To generate an installable "package" simply use cmake --install: For this step, your device needs to be configured for on-device development. Please see https://developer.android.com/studio/debug/dev-options for details. -Once ADB is enabled, use `adb push` to install `pkg-snapdragon` on the device. +Once ADB is enabled, use `adb push` to install `pkg-android` on the device. **Note that the toolchain Docker image doesn't have ADB and doesn't set up the ADB bridge. Please use native ADB on the host.** ``` -~/src/llama.cpp$ adb push pkg-snapdragon/llama.cpp /data/local/tmp/ -pkg-snapdragon/llama.cpp/bin/: 67 files pushed, 0 skipped. 190.2 MB/s (919095042 bytes in 4.607s) -pkg-snapdragon/llama.cpp/include/: 19 files pushed, 0 skipped. 20.5 MB/s (255173 bytes in 0.012s) -pkg-snapdragon/llama.cpp/lib/: 16 files pushed, 0 skipped. 144.4 MB/s (43801382 bytes in 0.289s) +~/src/llama.cpp$ adb push pkg-android/llama.cpp /data/local/tmp/ +pkg-android/llama.cpp/bin/: 67 files pushed, 0 skipped. 190.2 MB/s (919095042 bytes in 4.607s) +pkg-android/llama.cpp/include/: 19 files pushed, 0 skipped. 20.5 MB/s (255173 bytes in 0.012s) +pkg-android/llama.cpp/lib/: 16 files pushed, 0 skipped. 144.4 MB/s (43801382 bytes in 0.289s) 102 files pushed, 0 skipped. 186.9 MB/s (963151597 bytes in 4.914s) ``` @@ -115,24 +123,44 @@ Llama-3.2-1B-Instruct-Q4_0.gguf: 1 file pushed, 0 skipped. 38.3 MB/s (773025920 ### Windows -All artifacts are already installed in the `pkg-snapdragon` folder. -To run, adapt below instructions to use Powershell scripts in `scripts/snapdragon/windows`. +All artifacts are already installed in the `pkg-wos` folder. +To run, you can use the `scripts/snapdragon/run.py` runner script (see details below). ## How to Run -The easiest way to run llama.cpp cli tools is using provided wrapper scripts that properly set up all required environment variables. +The easiest way to run llama.cpp cli tools is using the provided `scripts/snapdragon/run.py` wrapper script. This script automatically +maps CLI options to environment variables, resolves executable paths, and runs the command locally, via ADB, or remotely via SSH on the +target device. -llama.cpp supports three backends on Snapdragon-based devices: CPU, Adreno GPU (GPUOpenCL), and Hexagon NPU (HTP0-4). -You can select which backend to run the model on using the `D=` variable, which maps to the `--device` option. +llama.cpp supports three backends on Snapdragon-based devices: CPU, Adreno GPU (GPUOpenCL), and Hexagon NPU. +You can select which backend(s) to run the model on using the `--device` option of the tool (or `--devices` option in `run.py`). Hexagon NPU behaves as a "GPU" device when it comes to `-ngl` and other offload-related options. -Here are some examples of running various llama.cpp tools via ADB. +Here are some examples of running various llama.cpp tools. -Simple question for Llama-3.2-1B +Generating a completion with Gemma on Android (relying on default `HTP0:0` device and default thread count `-t 6`): ``` -~/src/llama.cpp$ M=Llama-3.2-1B-Instruct-Q4_0.gguf D=HTP0 ./scripts/snapdragon/adb/run-completion.sh -p "what is the most popular cookie in the world?" +~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb -- llama-completion -m models/gemma-2-2b-it-Q4_0.gguf -f prompts/sample_prompt_1024.txt --jinja -st +... +ggml-hex: Hexagon backend (experimental) : allocating new registry : ndev 1 +ggml-hex: Hexagon Arch version v79 +ggml-hex: allocating new session: HTP0:0 +... +load_tensors: offloading output layer to GPU +load_tensors: offloaded 27/27 layers to GPU +load_tensors: CPU model buffer size = 300.00 MiB +load_tensors: HTP0:0 model buffer size = 1400.26 MiB +... +llama_perf_context_print: prompt eval time = 320.00 ms / 1024 tokens ( 0.31 ms per token, 3200.00 tokens per second) +llama_perf_context_print: eval time = 2100.00 ms / 100 runs ( 21.00 ms per token, 47.62 tokens per second) +``` + +Simple question for Llama-3.2-1B: + +``` +~/src/llama.cpp$ ./scripts/snapdragon/run.py --target android --devices HTP0 -- llama-cli -m Llama-3.2-1B-Instruct-Q4_0.gguf -p "what is the most popular cookie in the world?" ... ggml-hex: Hexagon backend (experimental) : allocating new registry : ndev 1 ggml-hex: Hexagon Arch version v79 @@ -142,8 +170,7 @@ ggml-hex: new session: HTP0 : session-id 0 domain-id 3 uri file:///libggml-htp-v load_tensors: offloading output layer to GPU load_tensors: offloaded 17/17 layers to GPU load_tensors: CPU model buffer size = 225.49 MiB -load_tensors: HTP0 model buffer size = 0.26 MiB -load_tensors: HTP0-REPACK model buffer size = 504.00 MiB +load_tensors: HTP0 model buffer size = 504.26 MiB ... I hope this helps you understand the world's most popular cookies! [end of text] ... @@ -156,60 +183,25 @@ llama_perf_context_print: graphs reused = 473 llama_memory_breakdown_print: | memory breakdown [MiB] | total free self model context compute unaccounted | llama_memory_breakdown_print: | - HTP0 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | llama_memory_breakdown_print: | - Host | 439 = 225 + 136 + 77 | -llama_memory_breakdown_print: | - HTP0-REPACK | 504 = 504 + 0 + 0 | -``` - -Summary request for OLMoE-1B-7B. This is a large model that requires two HTP sessions/devices - -``` -~/src/llama.cpp$ M=OLMoE-1B-7B-0125-Instruct-Q4_0.gguf NDEV=2 D=HTP0,HTP1 ./scripts/snapdragon/adb/run-completion.sh -f surfing.txt -... -ggml-hex: Hexagon backend (experimental) : allocating new registry : ndev 1 -ggml-hex: Hexagon Arch version v81 -ggml-hex: allocating new session: HTP0 -ggml-hex: allocating new session: HTP1 -... -load_tensors: offloading output layer to GPU -load_tensors: offloaded 17/17 layers to GPU -load_tensors: CPU model buffer size = 143.86 MiB -load_tensors: HTP1 model buffer size = 0.23 MiB -load_tensors: HTP1-REPACK model buffer size = 1575.00 MiB -load_tensors: HTP0 model buffer size = 0.28 MiB -load_tensors: HTP0-REPACK model buffer size = 2025.00 MiB -... -llama_context: CPU output buffer size = 0.19 MiB -llama_kv_cache: HTP1 KV buffer size = 238.00 MiB -llama_kv_cache: HTP0 KV buffer size = 306.00 MiB -llama_kv_cache: size = 544.00 MiB ( 8192 cells, 16 layers, 1/1 seqs), K (q8_0): 272.00 MiB, V (q8_0): 272.00 MiB -llama_context: HTP0 compute buffer size = 15.00 MiB -llama_context: HTP1 compute buffer size = 15.00 MiB -llama_context: CPU compute buffer size = 24.56 MiB -... -llama_perf_context_print: prompt eval time = 1730.57 ms / 212 tokens ( 8.16 ms per token, 122.50 tokens per second) -llama_perf_context_print: eval time = 5624.75 ms / 257 runs ( 21.89 ms per token, 45.69 tokens per second) -llama_perf_context_print: total time = 7377.33 ms / 469 tokens -llama_perf_context_print: graphs reused = 255 -llama_memory_breakdown_print: | memory breakdown [MiB] | total free self model context compute unaccounted | -llama_memory_breakdown_print: | - HTP0 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | -llama_memory_breakdown_print: | - HTP1 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | -llama_memory_breakdown_print: | - Host | 742 = 144 + 544 + 54 | -llama_memory_breakdown_print: | - HTP1-REPACK | 1575 = 1575 + 0 + 0 | -llama_memory_breakdown_print: | - HTP0-REPACK | 2025 = 2025 + 0 + 0 | ``` -Op test for MUL_MAT +Op test for MUL_MAT: ``` -~/src/llama.cpp$ HB=0 ./scripts/snapdragon/adb/run-tool.sh test-backend-ops -b HTP0 -o MUL_MAT +~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb --hex-hostbuf 0 --devices HTP0:0 -- test-backend-ops -b HTP0:0 -o MUL_MAT ... -Backend 2/3: HTP0 +Backend 2/3: HTP0:0 Device description: Hexagon Device memory: 2048 MB (2048 MB free) MUL_MAT(type_a=q4_0,type_b=f32,m=16,n=1,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK MUL_MAT(type_a=q4_0,type_b=f32,m=16,n=2,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK MUL_MAT(type_a=q4_0,type_b=f32,m=16,n=3,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK +``` + +Llama benchmark: -~/src/llama.cpp-hexagon$ M=Llama-3.2-1B-Instruct-Q4_0.gguf ./scripts/snapdragon/adb/run-bench.sh -p 128 -n 64 +``` +~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb --devices HTP0 -- llama-bench -p 128 -n 64 -m Llama-3.2-1B-Instruct-Q4_0.gguf ... ggml-hex: Hexagon backend (experimental) : allocating new registry : ndev 1 ggml-hex: Hexagon Arch version v79 @@ -219,15 +211,20 @@ ggml-hex: new session: HTP0 : session-id 0 domain-id 3 uri file:///libggml-htp-v | ---------------| ---------: | -----: | ---------- | --: | ------: | ------: | ---: | ----: | ------------: | | llama 1B Q4_0 | 729.75 MiB | 1.24 B | HTP | 99 | 4 | 128 | 0 | pp128 | 169.42 ± 1.75 | | llama 1B Q4_0 | 729.75 MiB | 1.24 B | HTP | 99 | 4 | 128 | 0 | tg64 | 51.54 ± 1.13 | - -build: 6a8cf8914 (6733) ``` ## Environment variables -- `GGML_HEXAGON_NDEV=1` - Controls the number of devices/sessions to allocate. The default is 1. - Most quantized models under 4B fit into a single session; an 8B model needs two, and a 20B model needs four. +- `GGML_HEXAGON_DEVICES` (default: not set, defaults to HTP0 session) + Controls which NPU devices and sessions to allocate. Can be configured as: + - A single integer `N`: Allocates `N` sessions named `HTP0`, `HTP1`, ..., `HTP` (behaves identically to `GGML_HEXAGON_NDEV=N`). + - A comma-separated list of device names in `HTP:` format (or legacy `HTP` format). For example, `HTP0:0,HTP0:1` creates two virtual + sessions on the first physical NPU (useful for memory limits). `HTP0:0,HTP1:0` allocates one session on each of the two physical NPUs + on a dual-NPU device. + +- `GGML_HEXAGON_NDEV` (deprecated) + Replaced by `GGML_HEXAGON_DEVICES`. Controls the number of virtual sessions to allocate on physical NPU `0`. + Allocates sessions named `HTP0`, `HTP1`, etc. - `GGML_HEXAGON_NHVX=0` Controls the number of HVX hardware threads to use. The default is all (actual number varies depending on the hardware version). @@ -255,26 +252,17 @@ build: 6a8cf8914 (6733) - `2` Extended profile with per-op `usecs`, `cycles` and default PMU counter data - `0x1,...,0x8` Extended profile with per-op `usecs`, `cycles` and custom PMU counter data - The logging output can be either saved into a file for post-processing or it can be piped directly into the post-processing tool to generate the report. + The logging output can be either saved into a file for post-processing or it can be piped directly into the post-processing tool + to generate the report. Examples: - `GGML_HEXAGON_PROFILE=1 llama-completion ... |& ./scripts/snapdragon/ggml-hexagon-profile.py -` - -- `GGML_HEXAGON_OPSTAGE=0x0` - Allows enabling specific stages of the Op processing pipeline: - - - `0x1` Enable Op Queue (i.e., queuing Ops into NPU) - - `0x2` Enable Op Compute (MUL_MAT, etc.) - - Examples: - - `GGML_HEXAGON_OPSTAGE=0x1 llama-completion ...` - Ops are enqueued to the NPU but dma & compute are disabled - `GGML_HEXAGON_OPSTAGE=0x3 llama-completion ...` - Full queuing and processing of Ops (default) + `GGML_HEXAGON_PROFILE=1 ./scripts/snapdragon/run.py --target adb -- llama-cli ... |& ./scripts/snapdragon/ggml-hexagon-profile.py -` - `GGML_HEXAGON_OPFILTER=regex` Allows filtering (disabling) Ops that match the regex pattern: Examples: - `GGML_HEXAGON_OPFILTER="FLASH_ATTN_EXT" llama-completion ...` - Disable Flash Attention on Hexagon (falls back to CPU or GPU) - `GGML_HEXAGON_OPFILTER="ADD\|SUB" llama-completion ...` - Disable ADD and SUB on Hexagon (fall back to CPU or GPU) + `GGML_HEXAGON_OPFILTER="FLASH_ATTN_EXT" ./scripts/snapdragon/run.py --target adb -- llama-cli ...` - Disable Flash Attention on Hexagon (falls back to CPU or GPU) + `GGML_HEXAGON_OPFILTER="ADD\|SUB" ./scripts/snapdragon/run.py --target adb -- llama-cli ...` - Disable ADD and SUB on Hexagon (fall back to CPU or GPU) + diff --git a/docs/backend/snapdragon/developer.md b/docs/backend/snapdragon/developer.md index 9d56638e3d5..d7d9f2a2790 100644 --- a/docs/backend/snapdragon/developer.md +++ b/docs/backend/snapdragon/developer.md @@ -39,22 +39,21 @@ the repacking. ## Large model handling -Hexagon NPU session (aka Process Domain (PD) in the Hexagon docs) is limited to a memory mapping of around 3.5GB. -In llama.cpp/GGML the Hexagon session is mapped to a single GGML backend device (HTP0, HTP1, etc). +Hexagon NPU sessions (aka Process Domains (PD) in the Hexagon SDK) are limited to a maximum memory mapping window of around 3.5GB. +In llama.cpp/GGML, each Hexagon session is mapped to a single GGML backend device (e.g., `HTP0:0`, `HTP0:1`, etc. when using +`GGML_HEXAGON_DEVICES`, or `HTP0`, `HTP1` in legacy mode). -In order to map models larger than 3.5GB we need to allocate multiple devices and split the model. -For this we're taking advantage of the llama.cpp/GGML multi-GPU layer-splitting support. -Each Hexagon device behaves like a GPU from the offload and model splitting perspective. +To support running models larger than 3.5GB on a single device, the Hexagon backend dynamically maps and unmaps execution buffers +during the graph execution cycle to stay within the Process Domain window. This enables large models to run successfully on a single +NPU device. -Here is an example of running GPT-OSS-20B model on a newer Snapdragon device with 16GB of DDR. +Alternatively, users can choose to use standard llama.cpp/GGML layer-splitting mode to partition and split the model across +multiple Hexagon devices or virtual sessions (which behave like multiple GPUs from the offload and splitting perspective). + +Here is an example of running GPT-OSS-20B model on a Snapdragon device using 4 virtual sessions on a single NPU (physical index 0). ``` -M=gpt-oss-20b-Q4_0.gguf NDEV=4 D=HTP0,HTP1,HTP2,HTP3 P=surfing.txt scripts/snapdragon/adb/run-completion.sh -f surfing.txt -n 32 -... -LD_LIBRARY_PATH=/data/local/tmp/llama.cpp/lib -ADSP_LIBRARY_PATH=/data/local/tmp/llama.cpp/lib -GGML_HEXAGON_NDEV=4 ./bin/llama-cli --load-mode none -m /data/local/tmp/llama.cpp/../gguf/gpt-oss-20b-Q4_0.gguf - -t 4 --ctx-size 8192 --batch-size 128 -ctk q8_0 -ctv q8_0 -fa on -ngl 99 --device HTP0,HTP1,HTP2,HTP3 -no-cnv -f surfing.txt +~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb --devices HTP0:0,HTP0:1,HTP0:2,HTP0:3 -- llama-cli --load-mode none -m /data/local/tmp/gguf/gpt-oss-20b-Q4_0.gguf -t 4 --ctx-size 8192 --batch-size 128 -ctk q8_0 -ctv q8_0 -fa on -ngl 99 -no-cnv -f surfing.txt ... llama_model_loader: - type f32: 289 tensors llama_model_loader: - type q4_0: 96 tensors @@ -63,33 +62,29 @@ llama_model_loader: - type mxfp4: 72 tensors ... load_tensors: offloaded 25/25 layers to GPU load_tensors: CPU model buffer size = 1182.09 MiB -load_tensors: HTP1 model buffer size = 6.64 MiB -load_tensors: HTP1-REPACK model buffer size = 2505.94 MiB -load_tensors: HTP3 model buffer size = 5.55 MiB -load_tensors: HTP3-REPACK model buffer size = 2088.28 MiB -load_tensors: HTP0 model buffer size = 7.75 MiB -load_tensors: HTP0-REPACK model buffer size = 2923.59 MiB -load_tensors: HTP2 model buffer size = 6.64 MiB -load_tensors: HTP2-REPACK model buffer size = 2505.94 MiB +load_tensors: HTP0:1 model buffer size = 2512.58 MiB +load_tensors: HTP0:3 model buffer size = 2093.83 MiB +load_tensors: HTP0:0 model buffer size = 2931.34 MiB +load_tensors: HTP0:2 model buffer size = 2512.58 MiB ... llama_context: n_ctx_per_seq (8192) < n_ctx_train (131072) -- the full capacity of the model will not be utilized llama_context: CPU output buffer size = 0.77 MiB llama_kv_cache_iswa: creating non-SWA KV cache, size = 8192 cells -llama_kv_cache: HTP1 KV buffer size = 25.50 MiB -llama_kv_cache: HTP3 KV buffer size = 25.50 MiB -llama_kv_cache: HTP0 KV buffer size = 25.50 MiB -llama_kv_cache: HTP2 KV buffer size = 25.50 MiB +llama_kv_cache: HTP0:1 KV buffer size = 25.50 MiB +llama_kv_cache: HTP0:3 KV buffer size = 25.50 MiB +llama_kv_cache: HTP0:0 KV buffer size = 25.50 MiB +llama_kv_cache: HTP0:2 KV buffer size = 25.50 MiB llama_kv_cache: size = 102.00 MiB ( 8192 cells, 12 layers, 1/1 seqs), K (q8_0): 51.00 MiB, V (q8_0): 51.00 MiB llama_kv_cache_iswa: creating SWA KV cache, size = 256 cells -llama_kv_cache: HTP1 KV buffer size = 0.80 MiB -llama_kv_cache: HTP3 KV buffer size = 0.53 MiB -llama_kv_cache: HTP0 KV buffer size = 1.06 MiB -llama_kv_cache: HTP2 KV buffer size = 0.80 MiB +llama_kv_cache: HTP0:1 KV buffer size = 0.80 MiB +llama_kv_cache: HTP0:3 KV buffer size = 0.53 MiB +llama_kv_cache: HTP0:0 KV buffer size = 1.06 MiB +llama_kv_cache: HTP0:2 KV buffer size = 0.80 MiB llama_kv_cache: size = 3.19 MiB ( 256 cells, 12 layers, 1/1 seqs), K (q8_0): 1.59 MiB, V (q8_0): 1.59 MiB -llama_context: HTP0 compute buffer size = 16.06 MiB -llama_context: HTP1 compute buffer size = 16.06 MiB -llama_context: HTP2 compute buffer size = 16.06 MiB -llama_context: HTP3 compute buffer size = 16.06 MiB +llama_context: HTP0:0 compute buffer size = 16.06 MiB +llama_context: HTP0:1 compute buffer size = 16.06 MiB +llama_context: HTP0:2 compute buffer size = 16.06 MiB +llama_context: HTP0:3 compute buffer size = 16.06 MiB llama_context: CPU compute buffer size = 98.19 MiB ... llama_perf_context_print: prompt eval time = 3843.67 ms / 197 tokens ( 19.51 ms per token, 51.25 tokens per second) @@ -97,13 +92,9 @@ llama_perf_context_print: eval time = 1686.13 ms / 31 runs ( 54.3 llama_perf_context_print: total time = 6266.30 ms / 228 tokens llama_perf_context_print: graphs reused = 30 llama_memory_breakdown_print: | memory breakdown [MiB] | total free self model context compute unaccounted | -llama_memory_breakdown_print: | - HTP0 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | -llama_memory_breakdown_print: | - HTP1 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | -llama_memory_breakdown_print: | - HTP2 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | -llama_memory_breakdown_print: | - HTP3 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | +llama_memory_breakdown_print: | - HTP0:0 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | +llama_memory_breakdown_print: | - HTP0:1 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | +llama_memory_breakdown_print: | - HTP0:2 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | +llama_memory_breakdown_print: | - HTP0:3 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | llama_memory_breakdown_print: | - Host | 1476 = 1208 + 105 + 162 | -llama_memory_breakdown_print: | - HTP1-REPACK | 2505 = 2505 + 0 + 0 | -llama_memory_breakdown_print: | - HTP3-REPACK | 2088 = 2088 + 0 + 0 | -llama_memory_breakdown_print: | - HTP0-REPACK | 2923 = 2923 + 0 + 0 | -llama_memory_breakdown_print: | - HTP2-REPACK | 2505 = 2505 + 0 + 0 | ``` diff --git a/docs/backend/snapdragon/linux.md b/docs/backend/snapdragon/linux.md index 90fdadb6c93..d4ecd9b1b16 100644 --- a/docs/backend/snapdragon/linux.md +++ b/docs/backend/snapdragon/linux.md @@ -1,25 +1,37 @@ # Snapdragon-based Linux devices -## Docker Setup +The cross-compilation is performed using the Snapdragon Linux Docker toolchain image (see +[github.com/snapdragon-toolchain](https://github.com/snapdragon-toolchain)): -The easiest way to build llama.cpp for a Snapdragon-based Linux device is using the toolchain Docker image (see [github.com/snapdragon-toolchain](https://github.com/snapdragon-toolchain)). -This image includes OpenCL SDK, Hexagon SDK, CMake, and the ARM64 Linux cross-compilation toolchain. +* **Linux toolchain**: `ghcr.io/snapdragon-toolchain/arm64-linux:v0.7` -Cross-compilation is supported on **Linux X86** hosts. The resulting binaries are deployed to and run on the target **Qualcomm Snapdragon ARM64 Linux** device. +The unified build utility (`scripts/snapdragon/build.py`) automatically pulls +and orchestrates this container to perform target compilation. You only need to +ensure that Docker is running on your host machine. + +## How to Build + +### Using build.py script (Recommended) + +The easiest way to build llama.cpp is by using the `scripts/snapdragon/build.py` script. It automatically copies the CMake presets, +launches the correct compilation Docker container, builds the libraries and tools, +installs them, and optionally pushes them to your target device. + +Build and deploy for a Linux target (using SSH deployment alias `lnx` or `linux`): ``` -~/src/llama.cpp$ docker run -it -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-linux:v0.1 -[d]/> cd /workspace +$ ./scripts/snapdragon/build.py --target lnx:user@host --push ``` -Note: The rest of the **Linux** build process assumes that you're running inside the toolchain container. +### Manual CMake Build +Alternatively, you can build llama.cpp manually by entering the cross-compilation Docker container and running the CMake commands: -## How to Build +```bash +# Start the cross-compilation container manually: +~/src/llama.cpp$ docker run -it --rm -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-linux:v0.7 -Let's build llama.cpp with CPU, OpenCL, and Hexagon backends via CMake presets: - -``` +# Inside the container, build the project using presets: [d]/workspace> cp docs/backend/snapdragon/CMakeUserPresets.json . [d]/workspace> cmake --preset arm64-linux-snapdragon-release -B build-snapdragon @@ -30,17 +42,19 @@ Let's build llama.cpp with CPU, OpenCL, and Hexagon backends via CMake presets: To generate an installable "package" simply use cmake --install, then zip it: ``` -[d]/workspace> cmake --install build-snapdragon --prefix pkg-snapdragon -[d]/workspace> zip -r pkg-snapdragon.zip pkg-snapdragon +[d]/workspace> cmake --install build-snapdragon --prefix pkg-linux +[d]/workspace> zip -r pkg-linux.zip pkg-linux ``` ## How to Install -For this step, you will deploy the built binaries and libraries to the target Linux device. Transfer `pkg-snapdragon.zip` to the target device, then unzip it and set up the environment variables: +For this step, you will deploy the built binaries and libraries to the target +Linux device. Transfer `pkg-linux.zip` to the target device, then unzip it +and set up the environment variables: ``` -$ unzip pkg-snapdragon.zip -$ cd pkg-snapdragon +$ unzip pkg-linux.zip +$ cd pkg-linux $ export LD_LIBRARY_PATH=./lib $ export ADSP_LIBRARY_PATH=./lib ``` @@ -52,7 +66,28 @@ $ wget https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/ ``` ## How to Run -Next, since we have setup the environment variables, we can run the llama-cli with the Hexagon backends: +You can run locally on the Snapdragon Linux device: +``` +$ ./scripts/snapdragon/run.py --devices HTP0 -- llama-cli -m Llama-3.2-3B-Instruct-Q4_0.gguf -ngl 99 -p "what is the most popular cookie in the world?" +``` + +Or run remotely from your host development machine using the SSH target option: +``` +$ ./scripts/snapdragon/run.py --target lnx:user@host --devices HTP0 -- llama-cli -m Llama-3.2-3B-Instruct-Q4_0.gguf -ngl 99 -p "what is the most popular cookie in the world?" +``` + +For multi-NPU systems, you can run a tensor split completion command targeting a remote Linux system: +``` +$ ./scripts/snapdragon/run.py --target ubuntu:maxk@192.168.1.87 --device HTP0:0,HTP1:0 -- llama-completion -m models/gemma-2b-it-Q4_0.gguf -f prompts/sample_prompt_1024.txt --jinja -st --split-mode tensor --ctx-size 8192 +``` + +This translates to the following command being executed remotely via SSH: +``` ++ ssh maxk@192.168.1.87 "cd ~/llama.cpp && ulimit -c unlimited && LD_LIBRARY_PATH=./lib ADSP_LIBRARY_PATH=./lib GGML_HEXAGON_DEVICES=HTP0:0,HTP1:0 GGML_HEXAGON_OPPOLL=1 ./bin/llama-completion -m models/gemma-2b-it-Q4_0.gguf -f prompts/sample_prompt_1024.txt --jinja -st --split-mode tensor --ctx-size 8192 -v -n 16 --device HTP0:0,HTP1:0 -ngl 99 --ubatch-size 1024 -fa on -t 6" +``` + +Alternatively, you can run the binary directly on the device: ``` $ ./bin/llama-cli -m Llama-3.2-3B-Instruct-Q4_0.gguf --device HTP0 -ngl 99 -p "what is the most popular cookie in the world?" ``` + diff --git a/docs/backend/snapdragon/windows.md b/docs/backend/snapdragon/windows.md index aa731413c90..886cfda3f65 100644 --- a/docs/backend/snapdragon/windows.md +++ b/docs/backend/snapdragon/windows.md @@ -1,3 +1,18 @@ +# Snapdragon-based Windows devices + +## Tool Dependencies + +Native Windows 11 arm64 builds have the following tool dependencies: +- MS Visual Studio 2026 (Community Edition or Pro) + - MSVC arm64 standard and runtime libraries + - UCRT and Driver Kit +- LLVM core libraries and Clang compiler (winget) +- CMake, Git, Python (winget) +- Hexagon SDK Community Edition 6.6 or later (see below) +- OpenCL SDK 2.3 or later (see below) + +Note: The rest of the **Windows** build process assumes that you're running natively in Powershell. + ## Overview The document covers procedures for installing the latest GPU and NPU drivers, and OpenCL and Hexagon SDKs. @@ -9,7 +24,18 @@ must be included in the .cat file digitally signed with a trusted certificate. This document covers details on how to generate personal certificate files (.pfx) and how to configure the system to allow for test signatures (aka test-signing). -## Install the latest Adreno OpenCL SDK +## Install Windows SDKs + +The recommended method is `setup-sdk.py`: + +``` +> python scripts\snapdragon\setup-sdk.py --list-sdk-releases +> python scripts\snapdragon\setup-sdk.py --hexagon --opencl +``` + +It installs the selected SDKs under `C:\Qualcomm` and sets their corresponding environment variables for the current user. Start a new terminal after it completes; native Windows builds check all SDK paths before CMake runs. + +Select the SDKs to install with `--hexagon` and `--opencl`; use both to prepare a dual-backend build. To select a different available version, pass it to the SDK option, for example `--hexagon 6.4.0.2`. SDK versions install side by side, so you can switch versions without deleting an existing installation. Use `--force` to reinstall the selected SDKs. Use a new CMake build directory after each switch because CMake caches the SDK paths. Either use the trimmed down version (optimized for CI) from @@ -53,7 +79,8 @@ Download the driver from https://softwarecenter.qualcomm.com/catalog/item/Qualcomm_HND -After the automated installation and reboot please make sure that the Hexagon NPU device shows up in the `Device Manager` (under `Neural Processors`). +After the automated installation and reboot please make sure that the Hexagon NPU device shows up in the `Device Manager` +(under `Neural Processors`). If the device is not available you can try installing all components (`qcnspmcdm8380`, `qcnspmcdm8380_ext`) manually. The components are extracted into @@ -130,12 +157,12 @@ However, additional settings are required for generating and signing HTP Ops lib > cmake --preset arm64-windows-snapdragon-release -B build-wos ... -> cmake --install build-wos --prefix pkg-snapdragon +> cmake --install build-wos --prefix pkg-wos ``` Once the build is complete HTP ops libraries will be installed like this ``` -> dir pkg-snapdragon/lib +> dir pkg-wos/lib ... -a---- 1/22/2026 6:01 PM 187656 libggml-htp-v73.so -a---- 1/22/2026 6:01 PM 191752 libggml-htp-v75.so @@ -147,8 +174,8 @@ Once the build is complete HTP ops libraries will be installed like this The .cat file, the signature and proper certificate installation can be verified with ``` -> signtool.exe verify /v /pa .\pkg-snapdragon\lib\libggml-htp.cat -Verifying: .\pkg-snapdragon\lib\libggml-htp.cat +> signtool.exe verify /v /pa .\pkg-wos\lib\libggml-htp.cat +Verifying: .\pkg-wos\lib\libggml-htp.cat Signature Index: 0 (Primary Signature) Hash of file (sha256): 9820C664DA59D5EAE31DBB664127FCDAEF59CDC31502496BC567544EC2F401CF @@ -156,6 +183,6 @@ Hash of file (sha256): 9820C664DA59D5EAE31DBB664127FCDAEF59CDC31502496BC567544EC Signing Certificate Chain: Issued to: GGML.HTP.v1 ... -Successfully verified: .\pkg-snapdragon\lib\libggml-htp.cat +Successfully verified: .\pkg-wos\lib\libggml-htp.cat ... ``` diff --git a/docs/ops.md b/docs/ops.md index 2c179dd01f3..54a3602833f 100644 --- a/docs/ops.md +++ b/docs/ops.md @@ -35,8 +35,8 @@ Legend: | COS | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | | COUNT_EQUAL | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | | CPY | ❌ | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ | -| CROSS_ENTROPY_LOSS | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | -| CROSS_ENTROPY_LOSS_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | +| CROSS_ENTROPY_LOSS | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | +| CROSS_ENTROPY_LOSS_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | | CUMSUM | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | | DIAG | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | | DIAG_MASK_INF | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ | diff --git a/docs/ops/Vulkan.csv b/docs/ops/Vulkan.csv index 59e67e1b208..13e67eb797f 100644 --- a/docs/ops/Vulkan.csv +++ b/docs/ops/Vulkan.csv @@ -19292,10 +19292,10 @@ "Vulkan0","FLASH_ATTN_EXT","hsk=128,hsv=128,nh=8,nr23=[4,1],kv=4096,nb=512,mask=1,sinks=0,max_bias=0.000000,logit_softcap=0.000000,prec=f32,type_K=f16,type_V=f16,permute=[0,1,2,3]","support","1","yes","Vulkan" "Vulkan0","FLASH_ATTN_EXT","hsk=256,hsv=256,nh=4,nr23=[6,1],kv=16384,nb=512,mask=1,sinks=0,max_bias=0.000000,logit_softcap=0.000000,prec=f32,type_K=f16,type_V=f16,permute=[0,1,2,3]","support","1","yes","Vulkan" "Vulkan0","FLASH_ATTN_EXT","hsk=128,hsv=128,nh=8,nr23=[4,1],kv=16384,nb=512,mask=1,sinks=0,max_bias=0.000000,logit_softcap=0.000000,prec=f32,type_K=f16,type_V=f16,permute=[0,1,2,3]","support","1","yes","Vulkan" -"Vulkan0","CROSS_ENTROPY_LOSS","type=f32,ne=[10,5,4,3]","support","0","no","Vulkan" -"Vulkan0","CROSS_ENTROPY_LOSS","type=f32,ne=[30000,1,1,1]","support","0","no","Vulkan" -"Vulkan0","CROSS_ENTROPY_LOSS_BACK","type=f32,ne=[10,5,4,3]","support","0","no","Vulkan" -"Vulkan0","CROSS_ENTROPY_LOSS_BACK","type=f32,ne=[30000,1,1,1]","support","0","no","Vulkan" +"Vulkan0","CROSS_ENTROPY_LOSS","type=f32,ne=[10,5,4,3]","support","1","yes","Vulkan" +"Vulkan0","CROSS_ENTROPY_LOSS","type=f32,ne=[30000,1,1,1]","support","1","yes","Vulkan" +"Vulkan0","CROSS_ENTROPY_LOSS_BACK","type=f32,ne=[10,5,4,3]","support","1","yes","Vulkan" +"Vulkan0","CROSS_ENTROPY_LOSS_BACK","type=f32,ne=[30000,1,1,1]","support","1","yes","Vulkan" "Vulkan0","OPT_STEP_ADAMW","type=f32,ne=[10,5,4,3]","support","1","yes","Vulkan" "Vulkan0","OPT_STEP_SGD","type=f32,ne=[10,5,4,3]","support","1","yes","Vulkan" "Vulkan0","GATED_DELTA_NET","type=f32,head_count=32,head_size=128,n_seq_tokens=1,n_seqs=1,v_repeat=1,permuted=0,kda=0,K=1","support","1","yes","Vulkan" diff --git a/docs/speculative.md b/docs/speculative.md index 0f9f8a3d977..ffb1e34c7fb 100644 --- a/docs/speculative.md +++ b/docs/speculative.md @@ -212,6 +212,15 @@ Use `--backend-sampling` to run supported target-model samplers on the model bac Unsupported samplers and device layouts fall back to CPU sampling. Tensor split mode does not support backend sampling. A fixed seed produces repeatable random draws, but stochastic CPU and backend sampling can still select different tokens because floating-point operations can differ between implementations and devices. Use greedy sampling when exact output matching is required. +### Synthetic Acceptance + +`llama-server` and `llama-cli` can replace normal speculative verification with synthetic decisions for benchmarking. The generated output is not valid model output because accepted draft tokens do not have to match the target model. + +Use exactly one of these options: + +- `--spec-synth-rates P0,P1,...` sets unconditional per-position acceptance probabilities. Entry `i` is the probability that the first `i+1` draft tokens are all accepted. The number of entries must match the effective maximum draft length. Values must be finite, within `[0, 1]`, and monotonically non-increasing. +- `--spec-synth-len L` sets the target mean acceptance length, including the target token. For `K` maximum draft tokens, `L` must be within `[1, K+1]`. The server finds a constant conditional probability `p` such that `p + p^2 + ... + p^K = L - 1`, then uses unconditional rates `[p, p^2, ..., p^K]`. + ### General Speculative Parameters ``` diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index 059e4496269..cbfe400139c 100644 --- a/ggml/include/ggml-rpc.h +++ b/ggml/include/ggml-rpc.h @@ -6,8 +6,8 @@ extern "C" { #endif -#define RPC_PROTO_MAJOR_VERSION 5 -#define RPC_PROTO_MINOR_VERSION 1 +#define RPC_PROTO_MAJOR_VERSION 6 +#define RPC_PROTO_MINOR_VERSION 0 #define RPC_PROTO_PATCH_VERSION 0 #ifdef __cplusplus diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 87ac0a702ef..b9c0fa3ddc0 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2936,12 +2936,13 @@ struct ggml_cplan ggml_graph_plan( const int64_t ne10 = node->src[1]->ne[0]; // W const int64_t ne11 = node->src[1]->ne[1]; // H const int64_t ne12 = node->src[1]->ne[2]; // Channels In + const int64_t ne13 = node->src[1]->ne[3]; // Batch GGML_ASSERT(node->src[0]->type == GGML_TYPE_F16 || node->src[0]->type == GGML_TYPE_F32); GGML_ASSERT(node->src[1]->type == GGML_TYPE_F32); cur += ggml_type_size(node->src[0]->type) * ne00 * ne01 * ne02 * ne03; - cur += ggml_type_size(node->src[0]->type) * ne10 * ne11 * ne12; + cur += ggml_type_size(node->src[0]->type) * ne10 * ne11 * ne12 * ne13; } break; case GGML_OP_TOP_K: diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index b869f4bddde..b47ce5463c6 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -7267,18 +7267,21 @@ static void ggml_compute_forward_conv_transpose_2d_impl( } } - // permute source data (src1) from (Sw x Sh x Cin) to (Cin x Sw x Sh) + // permute source data (src1) from (Sw x Sh x Cin) to (Cin x Sw x Sh), for all batches { kernel_t * const wdata = (kernel_t *) params->wdata + nk; - for (int i12 = 0; i12 < ne12; i12++) { - for (int i11 = 0; i11 < ne11; i11++) { - const float * const src = (float *)((char *) src1->data + i12*nb12 + i11*nb11); - kernel_t * dst_data = wdata + i11*ne10*ne12; - for (int i10 = 0; i10 < ne10; i10++) { - if constexpr (std::is_same_v) { - dst_data[i10*ne12 + i12] = GGML_CPU_FP32_TO_FP16(src[i10]); - } else { - dst_data[i10*ne12 + i12] = src[i10]; + for (int i13 = 0; i13 < ne13; i13++) { + kernel_t * const wdata_b = wdata + i13*ne10*ne11*ne12; + for (int i12 = 0; i12 < ne12; i12++) { + for (int i11 = 0; i11 < ne11; i11++) { + const float * const src = (float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11); + kernel_t * dst_data = wdata_b + i11*ne10*ne12; + for (int i10 = 0; i10 < ne10; i10++) { + if constexpr (std::is_same_v) { + dst_data[i10*ne12 + i12] = GGML_CPU_FP32_TO_FP16(src[i10]); + } else { + dst_data[i10*ne12 + i12] = src[i10]; + } } } } @@ -7305,24 +7308,27 @@ static void ggml_compute_forward_conv_transpose_2d_impl( kernel_t * const wdata_src = wdata + nk; for (int i2 = ip0; i2 < ip1; i2++) { // Cout - float * dst_data = (float *)((char *) dst->data + i2*nb2); kernel_t * wdata_kernel = wdata + i2*ne01*ne00*ne03; - for (int i11 = 0; i11 < ne11; i11++) { - for (int i10 = 0; i10 < ne10; i10++) { - const int i1n = i11*ne10*ne12 + i10*ne12; - for (int i01 = 0; i01 < ne01; i01++) { - for (int i00 = 0; i00 < ne00; i00++) { - float v = 0; - if constexpr (std::is_same_v) { - ggml_vec_dot_f16(ne03, &v, 0, - wdata_src + i1n, 0, - wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1); - } else { - ggml_vec_dot_f32(ne03, &v, 0, - wdata_src + i1n, 0, - wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1); + for (int i3 = 0; i3 < ne3; i3++) { // batch + float * dst_data = (float *)((char *) dst->data + i3*nb3 + i2*nb2); + kernel_t * wdata_src_b = wdata_src + i3*ne10*ne11*ne12; + for (int i11 = 0; i11 < ne11; i11++) { + for (int i10 = 0; i10 < ne10; i10++) { + const int i1n = i11*ne10*ne12 + i10*ne12; + for (int i01 = 0; i01 < ne01; i01++) { + for (int i00 = 0; i00 < ne00; i00++) { + float v = 0; + if constexpr (std::is_same_v) { + ggml_vec_dot_f16(ne03, &v, 0, + wdata_src_b + i1n, 0, + wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1); + } else { + ggml_vec_dot_f32(ne03, &v, 0, + wdata_src_b + i1n, 0, + wdata_kernel + i01*ne00*ne03 + i00*ne03, 0, 1); + } + dst_data[(i11*stride + i01)*ne0 + i10*stride + i00] += v; } - dst_data[(i11*stride + i01)*ne0 + i10*stride + i00] += v; } } } diff --git a/ggml/src/ggml-cuda/cpy.cu b/ggml/src/ggml-cuda/cpy.cu index fd7ffc0bc55..4d3cee4939f 100644 --- a/ggml/src/ggml-cuda/cpy.cu +++ b/ggml/src/ggml-cuda/cpy.cu @@ -386,46 +386,6 @@ static void ggml_cpy_f32_iq4_nl_cuda( (cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13); } -// check if a same-type copy reduces to a 2D strided copy (height rows of width -// contiguous bytes), so it can use cudaMemcpy2DAsync instead of the scalar kernel -static bool ggml_cuda_cpy_as_memcpy_2d(const ggml_tensor * src0, const ggml_tensor * src1, - size_t & width, size_t & height, size_t & spitch, size_t & dpitch) { - // require matching shape: a reshaped copy maps elements by flat order, which the - // prefix walk below does not handle - if (src0->type != src1->type || !ggml_are_same_shape(src0, src1)) { - return false; - } - - // grow the contiguous prefix block shared by both tensors - size_t block_nb = ggml_element_size(src0); - int d = 0; - for (; d < GGML_MAX_DIMS; ++d) { - if (src0->nb[d] != block_nb || src1->nb[d] != block_nb) { - break; - } - block_nb *= src0->ne[d]; - } - - // d == 0: nothing contiguous; d == GGML_MAX_DIMS: fully contiguous (handled by memcpy) - if (d == 0 || d == GGML_MAX_DIMS) { - return false; - } - - // dim d carries the rows; everything above it must be a single element - for (int i = d + 1; i < GGML_MAX_DIMS; ++i) { - if (src0->ne[i] != 1) { - return false; - } - } - - width = block_nb; - height = src0->ne[d]; - spitch = src0->nb[d]; - dpitch = src1->nb[d]; - - return spitch >= width && dpitch >= width; -} - void ggml_cuda_cpy(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, ggml_tensor * src1) { const int64_t ne = ggml_nelements(src0); GGML_ASSERT(ne == ggml_nelements(src1)); @@ -461,8 +421,6 @@ void ggml_cuda_cpy(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, gg const bool can_be_transposed = nb01 == (int64_t)ggml_element_size(src0) && src0->ne[3] == 1 && nb02 == ne00 * ne01 * (int64_t)ggml_element_size(src0); - size_t mc_width = 0, mc_height = 0, mc_spitch = 0, mc_dpitch = 0; - if (src0->type == src1->type && contiguous_srcs) { GGML_ASSERT(ggml_nbytes(src0) == ggml_nbytes(src1)); #if defined(GGML_USE_MUSA) && defined(GGML_MUSA_MUDNN_COPY) @@ -473,9 +431,6 @@ void ggml_cuda_cpy(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, gg { CUDA_CHECK(cudaMemcpyAsync(src1_ddc, src0_ddc, ggml_nbytes(src0), cudaMemcpyDeviceToDevice, main_stream)); } - } else if (ggml_cuda_cpy_as_memcpy_2d(src0, src1, mc_width, mc_height, mc_spitch, mc_dpitch)) { - CUDA_CHECK(cudaMemcpy2DAsync(src1_ddc, mc_dpitch, src0_ddc, mc_spitch, - mc_width, mc_height, cudaMemcpyDeviceToDevice, main_stream)); } else if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32) { if (can_be_transposed) { ggml_cpy_scalar_cuda diff --git a/ggml/src/ggml-cuda/mmq-config-pascal.cuh b/ggml/src/ggml-cuda/mmq-config-pascal-dp4a.cuh similarity index 99% rename from ggml/src/ggml-cuda/mmq-config-pascal.cuh rename to ggml/src/ggml-cuda/mmq-config-pascal-dp4a.cuh index e7d4a9a3fcb..83eb7c146e1 100644 --- a/ggml/src/ggml-cuda/mmq-config-pascal.cuh +++ b/ggml/src/ggml-cuda/mmq-config-pascal-dp4a.cuh @@ -1,4 +1,4 @@ -static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_pascal(ggml_type type, int J, bool fallback) { +static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_pascal_dp4a(ggml_type type, int J, bool fallback) { CASE(GGML_TYPE_Q1_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); CASE(GGML_TYPE_Q1_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); CASE(GGML_TYPE_Q1_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); diff --git a/ggml/src/ggml-cuda/mmq-config-pascal-older.cuh b/ggml/src/ggml-cuda/mmq-config-pascal-older.cuh new file mode 100644 index 00000000000..2a8dc9e1a93 --- /dev/null +++ b/ggml/src/ggml-cuda/mmq-config-pascal-older.cuh @@ -0,0 +1,273 @@ +static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config_pascal_older(ggml_type type, int J, bool fallback) { + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q1_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_1, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_1, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q2_K, 256, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q2_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q3_K, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 256, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 256, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 256, 1, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ1_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XXS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_XS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ2_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q3_K, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_S, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_XS, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ4_NL, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + +// --------------------------------------------------------------------------------------------- + + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_MXFP4, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 8, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 24, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 40, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_NVFP4, 256, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_NVFP4, MMQ_ITER_K, false, false); + + return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true); +} diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu index 707437ea3e5..7fb4401489c 100644 --- a/ggml/src/ggml-cuda/mmq.cu +++ b/ggml/src/ggml-cuda/mmq.cu @@ -314,7 +314,9 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t } if (ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_DP4A) { - return false; + // for MoE, mmq is faster even without native dp4a + // TODO: check if cards older than pascal might benefit from this as well + return cc >= GGML_CUDA_CC_PASCAL && n_experts > 0; } #ifdef GGML_CUDA_FORCE_MMQ diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh index 2eb15fdfad9..c978b4421c5 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -213,7 +213,8 @@ struct ggml_cuda_mmq_config { return ggml_cuda_mmq_config((type_), (nthreads_), (occupancy_), (I_), (J_), (sram_layout_), (K_vram_), (stream_k_), (fallback_)); \ } \ -#include "mmq-config-pascal.cuh" +#include "mmq-config-pascal-older.cuh" +#include "mmq-config-pascal-dp4a.cuh" #include "mmq-config-ampere.cuh" #include "mmq-config-blackwell.cuh" @@ -247,7 +248,10 @@ static __host__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(const ggml_type ty if (ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) { return ggml_cuda_mmq_get_config_ampere(type, J, fallback); } - return ggml_cuda_mmq_get_config_pascal(type, J, fallback); + if (ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_DP4A) { + return ggml_cuda_mmq_get_config_pascal_dp4a(type, J, fallback); + } + return ggml_cuda_mmq_get_config_pascal_older(type, J, fallback); } static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_type type, int J, bool fallback) { @@ -268,8 +272,10 @@ static constexpr __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_config(ggml_t return ggml_cuda_mmq_get_config_blackwell(type, J, fallback); #elif __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA return ggml_cuda_mmq_get_config_ampere(type, J, fallback); +#elif __CUDA_ARCH__ >= GGML_CUDA_CC_DP4A + return ggml_cuda_mmq_get_config_pascal_dp4a(type, J, fallback); #else - return ggml_cuda_mmq_get_config_pascal(type, J, fallback); + return ggml_cuda_mmq_get_config_pascal_older(type, J, fallback); #endif // BLACKWELL_MMA_AVAILABLE #endif // GGML_USE_HIP GGML_UNUSED_VARS(type, J, fallback); diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index e8a5009b381..53e86075591 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -6,6 +6,7 @@ #include #include +#include #include #include #include @@ -18,6 +19,7 @@ #include #include #include +#include #include #ifdef _WIN32 @@ -52,6 +54,8 @@ #include "htp/matmul-ops.h" #include "htp/flash-attn-ops.h" #include "htp/unary-ops.h" +#include "htp/get-rows-ops.h" +#include "htp/set-rows-ops.h" #include "htp_iface.h" #include "htp-drv.h" @@ -59,6 +63,36 @@ using intvec = std::vector; using uintvec = std::vector; using u32vec = std::vector; +#define GGML_HEXAGON_MAX_SESSIONS 16 + +#define GGML_HEXAGON_FENCE_BUFFER_SIZE 8192 +#define GGML_HEXAGON_FENCE_SLOT_SIZE 128 + +struct ggml_hexagon_device_config { + int physical_idx = 0; + int virtual_idx = 0; + std::string name; +}; + +static ggml_hexagon_device_config opt_device_configs[GGML_HEXAGON_MAX_SESSIONS]; + +static int get_domain_id(int physical_idx) { + switch (physical_idx) { + case 0: return 3; // CDSP0 (all devices) + case 1: return 4; // CDSP1 (IQ9, IQ10) + case 2: return 18; // CDSP2 (IQ10) + case 3: return 19; // CDSP3 (IQ10) + default: return CDSP_DOMAIN_ID + physical_idx; + } +} + +static std::string get_domain_name(int physical_idx) { + if (physical_idx == 0) { + return CDSP_DOMAIN_NAME; + } + return std::string("cdsp") + std::to_string(physical_idx); +} + static int opt_arch = 0; // autodetect static size_t opt_ndev = 1; static size_t opt_nhvx = 0; // use all @@ -68,20 +102,19 @@ static size_t opt_mbuf = 1ul * 1024 * 1024 * 1024; // max buffer size static int opt_etm = 0; static int opt_verbose = 0; static int opt_profile = 0; // profiling mode (0-disabled, 1-basic, 2-pmu) -static int opt_hostbuf = 1; // hostbuf ON by default +static bool opt_hostbuf = false; static int opt_mm_select = 3; // 3 = HMX -> Tiled -> Flat -> CPU, 2 = Tiled -> Flat -> CPU, 1 = Flat -> CPU static int opt_fa_select = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported) +static int opt_ar_select = 2; // 2 = fused ALLREDUCE+ADD (DMA, default), 1 = unfused ALLREDUCE (DMA), 0 = fallback to CPY+FENCE // Default PMU events, if profiling with PMU (mode=2) is enabled // See https://docs.qualcomm.com/doc/80-N2040-60/topic/pmu-events.html // https://docs.qualcomm.com/doc/80-N2040-61/topic/hvx-pmu-events.html static u32vec opt_pmu_evt { 0x3, 0x111, 0x100, 0x105, 0x240, 0x256, 0x7D, 0x8C }; -// Enable all stages by default -static int opt_opstage = HTP_OPSTAGE_QUEUE | HTP_OPSTAGE_COMPUTE; static int opt_opbatch = 1024; // max number of ops in a batch -static int opt_opqueue = 16; // max number of pending batches +static int opt_opqueue = 64; // max number of pending batches static int opt_optrace = 0; // trace buffer size per thread (0 means default) static int opt_oppoll = 0; // polling for batch completions static int opt_opfusion = 1; // enable/disable op fusion @@ -121,7 +154,7 @@ static void ggml_hexagon_dump_op_exec(const std::string &sess_name, const htp_op static void ggml_hexagon_dump_op_supp(const std::string &sess_name, const struct ggml_tensor * op, bool supp) { if (!opt_verbose) return; - htp_opformat fmt(htp_opformat(htp_opnode{const_cast(op), {}, HTP_OP_INVALID})); + htp_opformat fmt(htp_opformat(htp_opnode(HTP_OP_INVALID, const_cast(op)))); GGML_LOG_DEBUG("ggml-hex: %s supports-op %s|%s|%s|%s|%s|%s|%s\n", sess_name.c_str(), ggml_op_desc(op), fmt.names, fmt.dims, fmt.types, fmt.strides, fmt.buffs, supp ? "yes" : "no"); } @@ -144,6 +177,7 @@ static const char * htp_event_name(uint16_t id) { case HTP_TRACE_EVT_L2FLUSH: return "L2FLUSH"; case HTP_TRACE_EVT_INIT: return "INIT"; case HTP_TRACE_EVT_BUFF: return "BUFF"; + case HTP_TRACE_EVT_FENCE: return "FENCE"; default: return "UNKNOWN"; } } @@ -205,7 +239,12 @@ static void ggml_hexagon_dump_trace_events(const std::string & sess_name, const } } -// ** +enum ggml_hexagon_tensor_flags { + GGML_HEXAGON_TENSOR_REPACK = (1 << 0), + GGML_HEXAGON_TENSOR_WEIGHT = (1 << 1), + GGML_HEXAGON_TENSOR_FENCE = (1 << 2), + GGML_HEXAGON_TENSOR_FUSEABLE = (1 << 3), +}; static inline bool ggml_hexagon_is_repack_type(enum ggml_type type) { return type == GGML_TYPE_Q4_0 || type == GGML_TYPE_Q4_1 || @@ -227,6 +266,15 @@ static void ggml_hexagon_precompute_matmul_params( struct htp_mm_kernel_params * kparams ); +static void ggml_hexagon_precompute_fused_matmul_add_params( + const struct ggml_hexagon_session * sess, + const struct ggml_tensor * src0, + const struct ggml_tensor * src1, + const struct ggml_tensor * src2, + const struct ggml_tensor * dst, + struct htp_mm_kernel_params * kparams +); + static void ggml_hexagon_precompute_unary_params( const struct ggml_hexagon_session * sess, uint32_t op, @@ -236,25 +284,75 @@ static void ggml_hexagon_precompute_unary_params( struct htp_unary_kernel_params * kparams ); -static void ggml_hexagon_precompute_fused_qkv_params( +static void ggml_hexagon_precompute_get_rows_params( const struct ggml_hexagon_session * sess, const struct ggml_tensor * src0, const struct ggml_tensor * src1, - struct htp_mm_kernel_params * kparams + const struct ggml_tensor * dst, + struct htp_get_rows_kernel_params * kparams ); -static void ggml_hexagon_precompute_fused_ffn_params( +static void ggml_hexagon_precompute_set_rows_params( const struct ggml_hexagon_session * sess, const struct ggml_tensor * src0, const struct ggml_tensor * src1, + const struct ggml_tensor * dst, + struct htp_set_rows_kernel_params * kparams +); + +static void ggml_hexagon_precompute_fused_mmnx_params( + const struct ggml_hexagon_session * sess, + const struct ggml_tensor * src0, + const struct ggml_tensor * src1, + int32_t n_weights, struct htp_mm_kernel_params * kparams ); +static bool ggml_hexagon_precompute_allreduce_params( + const struct ggml_hexagon_session * sess, + const struct ggml_tensor * dst, + uint32_t rank, + uint32_t n_ranks, + bool has_add, + bool is_row_bcast, + struct htp_allreduce_kernel_params * kparams +); + +static bool mm_is_hmx_eligible(const ggml_tensor * t); +static bool is_mergeable_mul_mat(const ggml_tensor * t); +static bool is_mergeable_mul_mat_pair(const ggml_tensor * n1, const ggml_tensor * n2); + // ** backend sessions +struct ggml_hexagon_tensor_extra { + std::vector shadow_buf; + size_t shadow_size { 0 }; + uint32_t flags { 0 }; +}; + +static inline bool ggml_hexagon_tensor_is_fuseable(const struct ggml_tensor * t) { + if (!t || !t->extra) return false; + auto extra = (const struct ggml_hexagon_tensor_extra *) t->extra; + return (extra->flags & GGML_HEXAGON_TENSOR_FUSEABLE) != 0; +} + +struct htp_opnode; + struct ggml_hexagon_opbatch; struct ggml_hexagon_opqueue; -struct htp_opnode; +struct ggml_hexagon_shared_buffer; +struct ggml_hexagon_session; + +struct ggml_backend_hexagon_comm_context { + std::vector backends; + size_t n_backends = 0; + uint32_t fence_seq = 0; +}; + +struct ggml_hexagon_event { + ggml_hexagon_session * sess = nullptr; + uint64_t seq = 0; +}; struct ggml_hexagon_session { std::string name; @@ -264,6 +362,8 @@ struct ggml_hexagon_session { uint32_t domain_id; uint64_t queue_id; int dev_id; + int phys_idx; + int virt_idx; bool valid_session; bool valid_handle; bool valid_queue; @@ -273,20 +373,24 @@ struct ggml_hexagon_session { ggml_hexagon_opbatch* op_batch; ggml_hexagon_opqueue* op_queue; + std::unordered_map> cloned_buffers; + std::unordered_set sync_peers; + ggml_backend_buffer_type buffer_type = {}; - ggml_backend_buffer_type repack_buffer_type = {}; + ggml_backend_buffer_type host_buffer_type = {}; - uint32_t n_threads = 0; - uint32_t n_hvx = 0; - uint32_t n_hmx = 0; - uint64_t vtcm_size = 0; - size_t max_vmem = 0; + uint32_t n_threads = 0; + uint32_t n_hvx = 0; + uint32_t n_hmx = 0; + uint64_t vtcm_size = 0; + size_t max_vmem = 0; size_t max_bufsize = 0; + uint32_t fence_seq; - struct { - uint64_t uid = 0; - std::vector htp_nodes; - } cached_graph; + uint64_t cached_uid = 0; + std::vector cached_nodes; + + mutable std::unordered_set needs_repack; ggml_hexagon_session(int dev_id, ggml_backend_dev_t dev) noexcept(false); ~ggml_hexagon_session() noexcept(true); @@ -297,10 +401,31 @@ struct ggml_hexagon_session { void release() noexcept(true); void enqueue_op(const htp_opnode & node); - void flush(bool all = true); + void enqueue_cpy(const ggml_tensor * src, ggml_tensor * dst, const ggml_tensor * sync_tensor = nullptr, uint32_t fence_seq = 0); + void enqueue_fence(const ggml_tensor * sync_tensor, uint32_t fence_seq = 0); + void enqueue_allreduce(const ggml_tensor * dst, const std::vector & src_tensors, const std::vector & sync_tensors, uint32_t rank, uint32_t n_ranks, uint32_t fence_seq_entry = 0, uint32_t fence_seq_exit = 0); + void flush(bool all = true); void flush_pending(bool all = false); - void flush_batch(); + void flush_batch(size_t min_ops = 1); + + uint64_t record_event(); + void wait_event(uint64_t seq); + + bool clone_buffer(const ggml_hexagon_shared_buffer*); + + void add_sync_peer(ggml_hexagon_session * peer) { + sync_peers.insert(peer); + } + + void flush_sync_peers() { + if (sync_peers.empty()) return; + + for (auto * peer : sync_peers) { + peer->flush_batch(); + } + sync_peers.clear(); + } }; // ** backend buffers @@ -315,26 +440,68 @@ struct ggml_backend_hexagon_buffer_type_context { std::string name; }; +struct ggml_hexagon_rpcmem_block { + uint8_t * base = nullptr; + int fd = -1; + size_t size = 0; + + ggml_hexagon_rpcmem_block(size_t size) { + base = (uint8_t *) rpcmem_alloc2(RPCMEM_HEAP_ID_SYSTEM, RPCMEM_DEFAULT_FLAGS, size); + if (!base) { + throw std::runtime_error("ggml-hex: rpcmem_alloc failed"); + } + fd = rpcmem_to_fd(base); + if (fd < 0) { + rpcmem_free(base); + throw std::runtime_error("ggml-hex: rpcmem_to_fd failed"); + } + this->size = size; + } + + ~ggml_hexagon_rpcmem_block() { + if (base) { + rpcmem_free(base); + } + } +}; + struct ggml_hexagon_shared_buffer { - ggml_hexagon_session * sess; - uint8_t * base; - size_t size; - int fd; - bool mapped; - bool pinned; + ggml_hexagon_session * sess; + std::shared_ptr mem; + std::vector tensor_extra; + uint32_t fence_head = 0; + size_t fences_size = 0; + bool mapped; + bool pinned; + + const char * c_name() const { return sess->c_name(); } + uint8_t * base() const { return mem ? mem->base : nullptr; } + size_t size() const { return mem ? mem->size : 0; } + int fd() const { return mem ? mem->fd : -1; } + + uint8_t * alloc_fence() { + if (fences_size == 0) return nullptr; + int max_slots = fences_size / GGML_HEXAGON_FENCE_SLOT_SIZE; + uint32_t slot = (fence_head++) % max_slots; + + size_t guard_offset = size() - fences_size; + uint8_t * fence_ptr = base() + guard_offset + (size_t)slot * GGML_HEXAGON_FENCE_SLOT_SIZE; + return fence_ptr; + } void mmap() { + if (!this->mem) return; fastrpc_map_flags flags = this->pinned ? FASTRPC_MAP_FD : FASTRPC_MAP_FD_DELAYED; - int err = fastrpc_mmap(sess->domain_id, this->fd, (void *) this->base, 0, this->size, flags); + int err = fastrpc_mmap(sess->domain_id, fd(), (void *) base(), 0, size(), flags); if (err != 0) { GGML_LOG_ERROR("ggml-hex: %s buffer mapping failed : domain_id %d size %zu fd %d error 0x%08x\n", sess->c_name(), - sess->domain_id, this->size, this->fd, (unsigned) err); + sess->domain_id, size(), fd(), (unsigned) err); throw std::runtime_error("ggml-hex: fastrpc_mmap failed (see log for details)"); } HEX_VERBOSE("ggml-hex: %s mapped buffer: base %p size %zu fd %d pinned %u\n", - sess->c_name(), (void *) this->base, this->size, this->fd, pinned); + sess->c_name(), (void *) base(), size(), fd(), pinned); this->mapped = true; } @@ -342,66 +509,69 @@ struct ggml_hexagon_shared_buffer { void unmap() { if (!this->mapped) return; - if (!this->pinned) { + if (!this->pinned && mem) { // HTP might still hold a reference, tell it drop it - htp_iface_munmap(sess->handle, this->fd); + htp_iface_munmap(sess->handle, fd()); } - fastrpc_munmap(sess->domain_id, this->fd, (void *) this->base, this->size); + if (mem) { + fastrpc_munmap(sess->domain_id, fd(), (void *) base(), size()); + } HEX_VERBOSE("ggml-hex: %s unmapped buffer: base %p size %zu fd %d\n", sess->c_name(), - (void *) this->base, size, this->fd); + (void *) base(), size(), fd()); this->mapped = false; - this->fd = -1; } void alloc(size_t size) { - if (this->base) return; - - this->base = (uint8_t *) rpcmem_alloc2(RPCMEM_HEAP_ID_SYSTEM, RPCMEM_DEFAULT_FLAGS, size); - if (!this->base) { - GGML_LOG_ERROR("ggml-hex: %s failed to allocate buffer : size %zu\n", sess->c_name(), size); - throw std::runtime_error("ggml-hex: rpcmem_alloc failed (see log for details)"); - } + if (this->mem) return; - this->fd = rpcmem_to_fd(this->base); - if (this->fd < 0) { - GGML_LOG_ERROR("ggml-hex: %s failed to get FD for buffer %p\n", sess->c_name(), (void *) this->base); - throw std::runtime_error("ggml-hex: rpcmem_to_fd failed (see log for details)"); - } - this->size = size; + this->mem = std::make_shared(size); HEX_VERBOSE("ggml-hex: %s allocated buffer: base %p size %zu fd %d pinned %d\n", sess->c_name(), - (void *) this->base, this->size, this->fd, (int) pinned); + (void *) base(), this->size(), fd(), (int) pinned); mmap(); } void free() { - if (!this->base) return; - unmap(); - rpcmem_free(this->base); - - HEX_VERBOSE("ggml-hex: %s freed buffer: base %p size %zu fd %d\n", sess->c_name(), - (void *) this->base, size, this->fd); + // The memory is freed when the shared_ptr refcount drops to 0. + HEX_VERBOSE("ggml-hex: %s release ref on buffer: base %p size %zu fd %d\n", sess->c_name(), + (void *) base(), size(), fd()); + this->mem = nullptr; + } + + ggml_hexagon_shared_buffer(ggml_hexagon_session * sess, size_t size, bool pinned = false, size_t fence_size = 0) { + this->sess = sess; + this->mapped = false; + this->pinned = pinned; + this->fences_size = fence_size; + + // Size adjustment inside the buffer class + size_t guard_offset = (size + 4095) & ~4095; + size_t total_size = guard_offset; + if (fence_size > 0) { + total_size += 4096 + fence_size; + } - this->base = NULL; + alloc(total_size); } - ggml_hexagon_shared_buffer(ggml_hexagon_session * sess, size_t size, bool pinned = false) { - this->sess = sess; - this->size = 0; - this->base = nullptr; - this->fd = -1; - this->mapped = false; - this->pinned = pinned; - - alloc(size); + // Clone constructor for cross-session mapping + ggml_hexagon_shared_buffer(ggml_hexagon_session * sess, const ggml_hexagon_shared_buffer & other) { + this->sess = sess; + this->mem = other.mem; + this->mapped = false; + this->pinned = other.pinned; + this->fences_size = other.fences_size; } ~ggml_hexagon_shared_buffer() { free(); + for (auto * extra : tensor_extra) { + delete extra; + } } }; @@ -416,18 +586,25 @@ static void ggml_backend_hexagon_buffer_free_buffer(ggml_backend_buffer_t buffer static void * ggml_backend_hexagon_buffer_get_base(ggml_backend_buffer_t buffer) { auto sbuf = static_cast(buffer->context); - return sbuf->base; + return sbuf->base(); } static enum ggml_status ggml_backend_hexagon_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { auto sbuf = static_cast(buffer->context); auto sess = sbuf->sess; - HEX_VERBOSE("ggml-hex: %s init-tensor %s : base %p data %p nbytes %zu usage %d\n", sess->c_name(), - tensor->name, (void *) sbuf->base, tensor->data, ggml_nbytes(tensor), (int) buffer->usage); + HEX_VERBOSE("ggml-hex: %s init-tensor %s : base %p data %p nbytes %zu\n", sess->c_name(), + tensor->name, (void *) sbuf->base(), tensor->data, ggml_nbytes(tensor)); + + auto extra = new ggml_hexagon_tensor_extra(); + sbuf->tensor_extra.push_back(extra); - if (tensor->view_src != NULL && tensor->view_offs == 0) { - return GGML_STATUS_SUCCESS; // nothing to do for the view + tensor->extra = extra; + if (ggml_hexagon_is_repack_type(tensor->type)) { + if (sess->needs_repack.count(tensor)) { + extra->flags |= GGML_HEXAGON_TENSOR_REPACK; + sess->needs_repack.erase(tensor); + } } return GGML_STATUS_SUCCESS; @@ -499,7 +676,7 @@ static void pack_mxfp4_quants(block_mxfp4 * x, const uint8_t * qs, unsigned int } // repack q4_0 data into q4_0_tiled tensor -static void repack_q4_0_tiled(ggml_tensor * t, const void * data, size_t size) { +static void repack_q4_0_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { const block_q4_0 * src_matrix = (const block_q4_0 *) data; int64_t ne0 = t->ne[0]; int64_t ne1 = t->ne[1]; @@ -513,46 +690,49 @@ static void repack_q4_0_tiled(ggml_tensor * t, const void * data, size_t size) { const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q4_0; const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - const block_q4_0 * src_expert = src_matrix + (i3 * ne2 + i2) * (ne1 * (ne0 / 32)); - uint8_t * matrix_dst = (uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; + size_t slice_size = ne1 * ggml_row_size(t->type, ne0); + int64_t start_slice = offset / slice_size; + int64_t end_slice = (offset + size + slice_size - 1) / slice_size; + if (end_slice > ne2 * ne3) { + end_slice = ne2 * ne3; + } - for (int ct = 0; ct < n_col_tiles; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - uint8_t * tile_dst = matrix_dst + (ct * n_k_tiles + kt) * tile_size; + for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { + const block_q4_0 * src_slice = src_matrix + (slice_idx - start_slice) * (ne1 * (ne0 / 32)); + uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size; - uint8_t tile_quants[32][32]; - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - unpack_q4_0_quants(tile_quants[row], &src_expert[r * (ne0 / 32) + kt], 0); - } else { - memset(tile_quants[row], 8, 32); - } - } + for (int ct = 0; ct < n_col_tiles; ct++) { + for (int kt = 0; kt < n_k_tiles; kt++) { + uint8_t * tile_dst = matrix_dst + (ct * n_k_tiles + kt) * tile_size; - for (int cp = 0; cp < 16; cp++) { - for (int row = 0; row < 32; row++) { - tile_dst[cp * 32 + row] = (tile_quants[row][2 * cp + 1] << 4) | tile_quants[row][2 * cp]; - } + uint8_t tile_quants[32][32]; + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r < ne1 && kt < ne0 / 32) { + unpack_q4_0_quants(tile_quants[row], &src_slice[r * (ne0 / 32) + kt], 0); + } else { + memset(tile_quants[row], 8, 32); } + } - ggml_half * scale_dst = (ggml_half *)(tile_dst + 512); + for (int cp = 0; cp < 16; cp++) { for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - scale_dst[row] = (r < ne1 && kt < ne0 / 32) ? src_expert[r * (ne0 / 32) + kt].d : 0; + tile_dst[cp * 32 + row] = (tile_quants[row][2 * cp + 1] << 4) | tile_quants[row][2 * cp]; } } + + ggml_half * scale_dst = (ggml_half *)(tile_dst + 512); + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + scale_dst[row] = (r < ne1 && kt < ne0 / 32) ? src_slice[r * (ne0 / 32) + kt].d : 0; + } } } } - - GGML_UNUSED(size); } // repack q4_0_tiled tensor into q4_0 data -static void repack_tiled_q4_0(void * data, const ggml_tensor * t, size_t size) { +static void repack_tiled_q4_0(void * data, const ggml_tensor * t, size_t offset, size_t size) { block_q4_0 * dst_matrix = (block_q4_0 *) data; int64_t ne0 = t->ne[0]; int64_t ne1 = t->ne[1]; @@ -566,48 +746,65 @@ static void repack_tiled_q4_0(void * data, const ggml_tensor * t, size_t size) { const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q4_0; const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - block_q4_0 * dst_expert = dst_matrix + (i3 * ne2 + i2) * (ne1 * (ne0 / 32)); - const uint8_t * matrix_src = (const uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; - - for (int ct = 0; ct < n_col_tiles; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - const uint8_t * tile_src = matrix_src + (ct * n_k_tiles + kt) * tile_size; - - uint8_t tile_quants[32][32]; - for (int cp = 0; cp < 16; cp++) { - for (int row = 0; row < 32; row++) { - uint8_t val = tile_src[cp * 32 + row]; - tile_quants[row][2 * cp + 0] = val & 0x0F; - tile_quants[row][2 * cp + 1] = val >> 4; - } - } + size_t slice_size = ne1 * ggml_row_size(t->type, ne0); + size_t row_size_bytes = ggml_row_size(t->type, ne0); + int64_t start_slice = offset / slice_size; + int64_t end_slice = (offset + size + slice_size - 1) / slice_size; + if (end_slice > ne2 * ne3) { + end_slice = ne2 * ne3; + } + + for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { + size_t cur_start_byte = (std::max)(offset, (size_t) slice_idx * slice_size); + size_t cur_end_byte = (std::min)(offset + size, (size_t) (slice_idx + 1) * slice_size); + size_t slice_offset_start = cur_start_byte - (size_t) slice_idx * slice_size; + size_t slice_offset_end = cur_end_byte - (size_t) slice_idx * slice_size; + + int64_t start_row = slice_offset_start / row_size_bytes; + int64_t end_row = (slice_offset_end + row_size_bytes - 1) / row_size_bytes; + end_row = (std::min)(end_row, ne1); + + int start_ct = start_row / 32; + int end_ct = (end_row + 31) / 32; + end_ct = (std::min)(end_ct, n_col_tiles); + + block_q4_0 * dst_slice = dst_matrix + (cur_start_byte - offset) / sizeof(block_q4_0); + const uint8_t * matrix_src = (const uint8_t *) t->data + slice_idx * matrix_size; + for (int ct = start_ct; ct < end_ct; ct++) { + for (int kt = 0; kt < n_k_tiles; kt++) { + const uint8_t * tile_src = matrix_src + (ct * n_k_tiles + kt) * tile_size; + + uint8_t tile_quants[32][32]; + for (int cp = 0; cp < 16; cp++) { for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - pack_q4_0_quants(&dst_expert[r * (ne0 / 32) + kt], tile_quants[row], 0); - } + uint8_t val = tile_src[cp * 32 + row]; + tile_quants[row][2 * cp + 0] = val & 0x0F; + tile_quants[row][2 * cp + 1] = val >> 4; } + } - const ggml_half * scale_src = (const ggml_half *)(tile_src + 512); - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - dst_expert[r * (ne0 / 32) + kt].d = scale_src[row]; - } + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r >= start_row && r < end_row && kt < ne0 / 32) { + pack_q4_0_quants(&dst_slice[(r - start_row) * (ne0 / 32) + kt], tile_quants[row], 0); + } + } + + const ggml_half * scale_src = (const ggml_half *)(tile_src + 512); + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r >= start_row && r < end_row && kt < ne0 / 32) { + dst_slice[(r - start_row) * (ne0 / 32) + kt].d = scale_src[row]; } } } } } - - GGML_UNUSED(size); } // repack q4_1 data into q4_1_tiled tensor -static void repack_q4_1_tiled(ggml_tensor * t, const void * data, size_t size) { +static void repack_q4_1_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { const block_q4_1 * src_matrix = (const block_q4_1 *) data; int64_t ne0 = t->ne[0]; int64_t ne1 = t->ne[1]; @@ -621,52 +818,55 @@ static void repack_q4_1_tiled(ggml_tensor * t, const void * data, size_t size) { const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q4_1; const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - const block_q4_1 * src_expert = src_matrix + (i3 * ne2 + i2) * (ne1 * (ne0 / 32)); - uint8_t * matrix_dst = (uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; + size_t slice_size = ne1 * ggml_row_size(t->type, ne0); + int64_t start_slice = offset / slice_size; + int64_t end_slice = (offset + size + slice_size - 1) / slice_size; + if (end_slice > ne2 * ne3) { + end_slice = ne2 * ne3; + } - for (int ct = 0; ct < n_col_tiles; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - uint8_t * tile_dst = matrix_dst + (ct * n_k_tiles + kt) * tile_size; + for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { + const block_q4_1 * src_slice = src_matrix + (slice_idx - start_slice) * (ne1 * (ne0 / 32)); + uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size; - uint8_t tile_quants[32][32]; - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - unpack_q4_1_quants(tile_quants[row], &src_expert[r * (ne0 / 32) + kt], 0); - } else { - memset(tile_quants[row], 0, 32); - } - } + for (int ct = 0; ct < n_col_tiles; ct++) { + for (int kt = 0; kt < n_k_tiles; kt++) { + uint8_t * tile_dst = matrix_dst + (ct * n_k_tiles + kt) * tile_size; - for (int cp = 0; cp < 16; cp++) { - for (int row = 0; row < 32; row++) { - tile_dst[cp * 32 + row] = (tile_quants[row][2 * cp + 1] << 4) | tile_quants[row][2 * cp]; - } + uint8_t tile_quants[32][32]; + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r < ne1 && kt < ne0 / 32) { + unpack_q4_1_quants(tile_quants[row], &src_slice[r * (ne0 / 32) + kt], 0); + } else { + memset(tile_quants[row], 0, 32); } + } - ggml_half * scale_dst = (ggml_half *)(tile_dst + 512); + for (int cp = 0; cp < 16; cp++) { for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - scale_dst[2 * row + 0] = src_expert[r * (ne0 / 32) + kt].d; - scale_dst[2 * row + 1] = src_expert[r * (ne0 / 32) + kt].m; - } else { - scale_dst[2 * row + 0] = 0; - scale_dst[2 * row + 1] = 0; - } + tile_dst[cp * 32 + row] = (tile_quants[row][2 * cp + 1] << 4) | tile_quants[row][2 * cp]; + } + } + + ggml_half * scale_dst = (ggml_half *)(tile_dst + 512); + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r < ne1 && kt < ne0 / 32) { + scale_dst[2 * row + 0] = src_slice[r * (ne0 / 32) + kt].d; + scale_dst[2 * row + 1] = src_slice[r * (ne0 / 32) + kt].m; + } else { + scale_dst[2 * row + 0] = 0; + scale_dst[2 * row + 1] = 0; } } } } } - - GGML_UNUSED(size); } // repack q4_1_tiled tensor into q4_1 data -static void repack_tiled_q4_1(void * data, const ggml_tensor * t, size_t size) { +static void repack_tiled_q4_1(void * data, const ggml_tensor * t, size_t offset, size_t size) { block_q4_1 * dst_matrix = (block_q4_1 *) data; int64_t ne0 = t->ne[0]; int64_t ne1 = t->ne[1]; @@ -680,49 +880,66 @@ static void repack_tiled_q4_1(void * data, const ggml_tensor * t, size_t size) { const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q4_1; const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - block_q4_1 * dst_expert = dst_matrix + (i3 * ne2 + i2) * (ne1 * (ne0 / 32)); - const uint8_t * matrix_src = (const uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; - - for (int ct = 0; ct < n_col_tiles; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - const uint8_t * tile_src = matrix_src + (ct * n_k_tiles + kt) * tile_size; - - uint8_t tile_quants[32][32]; - for (int cp = 0; cp < 16; cp++) { - for (int row = 0; row < 32; row++) { - uint8_t val = tile_src[cp * 32 + row]; - tile_quants[row][2 * cp + 0] = val & 0x0F; - tile_quants[row][2 * cp + 1] = val >> 4; - } - } + size_t slice_size = ne1 * ggml_row_size(t->type, ne0); + size_t row_size_bytes = ggml_row_size(t->type, ne0); + int64_t start_slice = offset / slice_size; + int64_t end_slice = (offset + size + slice_size - 1) / slice_size; + if (end_slice > ne2 * ne3) { + end_slice = ne2 * ne3; + } + for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { + size_t cur_start_byte = (std::max)(offset, (size_t) slice_idx * slice_size); + size_t cur_end_byte = (std::min)(offset + size, (size_t) (slice_idx + 1) * slice_size); + size_t slice_offset_start = cur_start_byte - (size_t) slice_idx * slice_size; + size_t slice_offset_end = cur_end_byte - (size_t) slice_idx * slice_size; + + int64_t start_row = slice_offset_start / row_size_bytes; + int64_t end_row = (slice_offset_end + row_size_bytes - 1) / row_size_bytes; + end_row = (std::min)(end_row, ne1); + + int start_ct = start_row / 32; + int end_ct = (end_row + 31) / 32; + end_ct = (std::min)(end_ct, n_col_tiles); + + block_q4_1 * dst_slice = dst_matrix + (cur_start_byte - offset) / sizeof(block_q4_1); + const uint8_t * matrix_src = (const uint8_t *) t->data + slice_idx * matrix_size; + + for (int ct = start_ct; ct < end_ct; ct++) { + for (int kt = 0; kt < n_k_tiles; kt++) { + const uint8_t * tile_src = matrix_src + (ct * n_k_tiles + kt) * tile_size; + + uint8_t tile_quants[32][32]; + for (int cp = 0; cp < 16; cp++) { for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - pack_q4_1_quants(&dst_expert[r * (ne0 / 32) + kt], tile_quants[row], 0); - } + uint8_t val = tile_src[cp * 32 + row]; + tile_quants[row][2 * cp + 0] = val & 0x0F; + tile_quants[row][2 * cp + 1] = val >> 4; } + } - const ggml_half * scale_src = (const ggml_half *)(tile_src + 512); - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - dst_expert[r * (ne0 / 32) + kt].d = scale_src[2 * row]; - dst_expert[r * (ne0 / 32) + kt].m = scale_src[2 * row + 1]; - } + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r >= start_row && r < end_row && kt < ne0 / 32) { + pack_q4_1_quants(&dst_slice[(r - start_row) * (ne0 / 32) + kt], tile_quants[row], 0); + } + } + + const ggml_half * scale_src = (const ggml_half *)(tile_src + 512); + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r >= start_row && r < end_row && kt < ne0 / 32) { + dst_slice[(r - start_row) * (ne0 / 32) + kt].d = scale_src[2 * row]; + dst_slice[(r - start_row) * (ne0 / 32) + kt].m = scale_src[2 * row + 1]; } } } } } - - GGML_UNUSED(size); } // repack q8_0 data into q8_0_tiled tensor -static void repack_q8_0_tiled(ggml_tensor * t, const void * data, size_t size) { +static void repack_q8_0_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { const block_q8_0 * src_matrix = (const block_q8_0 *) data; int64_t ne0 = t->ne[0]; int64_t ne1 = t->ne[1]; @@ -736,41 +953,44 @@ static void repack_q8_0_tiled(ggml_tensor * t, const void * data, size_t size) { const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q8_0; const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - const block_q8_0 * src_expert = src_matrix + (i3 * ne2 + i2) * (ne1 * (ne0 / 32)); - uint8_t * matrix_dst = (uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; - - for (int ct = 0; ct < n_col_tiles; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - uint8_t * tile_dst = matrix_dst + (ct * n_k_tiles + kt) * tile_size; - - for (int cp = 0; cp < 16; cp++) { - int col0 = cp * 2; - int col1 = col0 + 1; - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - const block_q8_0 * b = (r < ne1 && kt < ne0 / 32) ? &src_expert[r * (ne0 / 32) + kt] : NULL; - tile_dst[cp * 64 + 2 * row + 0] = b ? b->qs[col0] : 0; - tile_dst[cp * 64 + 2 * row + 1] = b ? b->qs[col1] : 0; - } - } + size_t slice_size = ne1 * ggml_row_size(t->type, ne0); + int64_t start_slice = offset / slice_size; + int64_t end_slice = (offset + size + slice_size - 1) / slice_size; + if (end_slice > ne2 * ne3) { + end_slice = ne2 * ne3; + } + + for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { + const block_q8_0 * src_slice = src_matrix + (slice_idx - start_slice) * (ne1 * (ne0 / 32)); + uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size; - ggml_half * scale_dst = (ggml_half *)(tile_dst + 1024); + for (int ct = 0; ct < n_col_tiles; ct++) { + for (int kt = 0; kt < n_k_tiles; kt++) { + uint8_t * tile_dst = matrix_dst + (ct * n_k_tiles + kt) * tile_size; + + for (int cp = 0; cp < 16; cp++) { + int col0 = cp * 2; + int col1 = col0 + 1; for (int row = 0; row < 32; row++) { int64_t r = ct * 32 + row; - scale_dst[row] = (r < ne1 && kt < ne0 / 32) ? src_expert[r * (ne0 / 32) + kt].d : 0; + const block_q8_0 * b = (r < ne1 && kt < ne0 / 32) ? &src_slice[r * (ne0 / 32) + kt] : NULL; + tile_dst[cp * 64 + 2 * row + 0] = b ? b->qs[col0] : 0; + tile_dst[cp * 64 + 2 * row + 1] = b ? b->qs[col1] : 0; } } + + ggml_half * scale_dst = (ggml_half *)(tile_dst + 1024); + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + scale_dst[row] = (r < ne1 && kt < ne0 / 32) ? src_slice[r * (ne0 / 32) + kt].d : 0; + } } } } - - GGML_UNUSED(size); } // repack q8_0_tiled tensor into q8_0 data -static void repack_tiled_q8_0(void * data, const ggml_tensor * t, size_t size) { +static void repack_tiled_q8_0(void * data, const ggml_tensor * t, size_t offset, size_t size) { block_q8_0 * dst_matrix = (block_q8_0 *) data; int64_t ne0 = t->ne[0]; int64_t ne1 = t->ne[1]; @@ -784,45 +1004,62 @@ static void repack_tiled_q8_0(void * data, const ggml_tensor * t, size_t size) { const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q8_0; const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - block_q8_0 * dst_expert = dst_matrix + (i3 * ne2 + i2) * (ne1 * (ne0 / 32)); - const uint8_t * matrix_src = (const uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; - - for (int ct = 0; ct < n_col_tiles; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - const uint8_t * tile_src = matrix_src + (ct * n_k_tiles + kt) * tile_size; - - for (int cp = 0; cp < 16; cp++) { - int col0 = cp * 2; - int col1 = col0 + 1; - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - block_q8_0 & b = dst_expert[r * (ne0 / 32) + kt]; - b.qs[col0] = tile_src[cp * 64 + 2 * row + 0]; - b.qs[col1] = tile_src[cp * 64 + 2 * row + 1]; - } - } - } + size_t slice_size = ne1 * ggml_row_size(t->type, ne0); + size_t row_size_bytes = ggml_row_size(t->type, ne0); + int64_t start_slice = offset / slice_size; + int64_t end_slice = (offset + size + slice_size - 1) / slice_size; + if (end_slice > ne2 * ne3) { + end_slice = ne2 * ne3; + } + + for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { + size_t cur_start_byte = (std::max)(offset, (size_t) slice_idx * slice_size); + size_t cur_end_byte = (std::min)(offset + size, (size_t) (slice_idx + 1) * slice_size); + size_t slice_offset_start = cur_start_byte - (size_t) slice_idx * slice_size; + size_t slice_offset_end = cur_end_byte - (size_t) slice_idx * slice_size; + + int64_t start_row = slice_offset_start / row_size_bytes; + int64_t end_row = (slice_offset_end + row_size_bytes - 1) / row_size_bytes; + end_row = (std::min)(end_row, ne1); + + int start_ct = start_row / 32; + int end_ct = (end_row + 31) / 32; + end_ct = (std::min)(end_ct, n_col_tiles); + + block_q8_0 * dst_slice = dst_matrix + (cur_start_byte - offset) / sizeof(block_q8_0); + const uint8_t * matrix_src = (const uint8_t *) t->data + slice_idx * matrix_size; - const ggml_half * scale_src = (const ggml_half *)(tile_src + 1024); + for (int ct = start_ct; ct < end_ct; ct++) { + for (int kt = 0; kt < n_k_tiles; kt++) { + const uint8_t * tile_src = matrix_src + (ct * n_k_tiles + kt) * tile_size; + + for (int cp = 0; cp < 16; cp++) { + int col0 = cp * 2; + int col1 = col0 + 1; for (int row = 0; row < 32; row++) { int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - dst_expert[r * (ne0 / 32) + kt].d = scale_src[row]; + if (r >= start_row && r < end_row && kt < ne0 / 32) { + block_q8_0 & b = dst_slice[(r - start_row) * (ne0 / 32) + kt]; + b.qs[col0] = tile_src[cp * 64 + 2 * row + 0]; + b.qs[col1] = tile_src[cp * 64 + 2 * row + 1]; } } } + + const ggml_half * scale_src = (const ggml_half *)(tile_src + 1024); + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r >= start_row && r < end_row && kt < ne0 / 32) { + dst_slice[(r - start_row) * (ne0 / 32) + kt].d = scale_src[row]; + } + } } } } - - GGML_UNUSED(size); } // repack mxfp4 data into mxfp4_tiled tensor -static void repack_mxfp4_tiled(ggml_tensor * t, const void * data, size_t size) { +static void repack_mxfp4_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { const block_mxfp4 * src_matrix = (const block_mxfp4 *) data; int64_t ne0 = t->ne[0]; int64_t ne1 = t->ne[1]; @@ -836,46 +1073,49 @@ static void repack_mxfp4_tiled(ggml_tensor * t, const void * data, size_t size) const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_MXFP4; const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - const block_mxfp4 * src_expert = src_matrix + (i3 * ne2 + i2) * (ne1 * (ne0 / 32)); - uint8_t * matrix_dst = (uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; + size_t slice_size = ne1 * ggml_row_size(t->type, ne0); + int64_t start_slice = offset / slice_size; + int64_t end_slice = (offset + size + slice_size - 1) / slice_size; + if (end_slice > ne2 * ne3) { + end_slice = ne2 * ne3; + } - for (int ct = 0; ct < n_col_tiles; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - uint8_t * tile_dst = matrix_dst + (ct * n_k_tiles + kt) * tile_size; + for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { + const block_mxfp4 * src_slice = src_matrix + (slice_idx - start_slice) * (ne1 * (ne0 / 32)); + uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size; - uint8_t tile_quants[32][32]; - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - unpack_mxfp4_quants(tile_quants[row], &src_expert[r * (ne0 / 32) + kt], 0); - } else { - memset(tile_quants[row], 0, 32); - } - } + for (int ct = 0; ct < n_col_tiles; ct++) { + for (int kt = 0; kt < n_k_tiles; kt++) { + uint8_t * tile_dst = matrix_dst + (ct * n_k_tiles + kt) * tile_size; - for (int cp = 0; cp < 16; cp++) { - for (int row = 0; row < 32; row++) { - tile_dst[cp * 32 + row] = (tile_quants[row][2 * cp + 1] << 4) | tile_quants[row][2 * cp]; - } + uint8_t tile_quants[32][32]; + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r < ne1 && kt < ne0 / 32) { + unpack_mxfp4_quants(tile_quants[row], &src_slice[r * (ne0 / 32) + kt], 0); + } else { + memset(tile_quants[row], 0, 32); } + } - uint8_t * scale_dst = tile_dst + 512; + for (int cp = 0; cp < 16; cp++) { for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - scale_dst[row] = (r < ne1 && kt < ne0 / 32) ? src_expert[r * (ne0 / 32) + kt].e : 0; + tile_dst[cp * 32 + row] = (tile_quants[row][2 * cp + 1] << 4) | tile_quants[row][2 * cp]; } } + + uint8_t * scale_dst = tile_dst + 512; + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + scale_dst[row] = (r < ne1 && kt < ne0 / 32) ? src_slice[r * (ne0 / 32) + kt].e : 0; + } } } } - - GGML_UNUSED(size); } // repack mxfp4_tiled tensor into mxfp4 data -static void repack_tiled_mxfp4(void * data, const ggml_tensor * t, size_t size) { +static void repack_tiled_mxfp4(void * data, const ggml_tensor * t, size_t offset, size_t size) { block_mxfp4 * dst_matrix = (block_mxfp4 *) data; int64_t ne0 = t->ne[0]; int64_t ne1 = t->ne[1]; @@ -889,133 +1129,179 @@ static void repack_tiled_mxfp4(void * data, const ggml_tensor * t, size_t size) const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_MXFP4; const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - block_mxfp4 * dst_expert = dst_matrix + (i3 * ne2 + i2) * (ne1 * (ne0 / 32)); - const uint8_t * matrix_src = (const uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; - - for (int ct = 0; ct < n_col_tiles; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - const uint8_t * tile_src = matrix_src + (ct * n_k_tiles + kt) * tile_size; - - uint8_t tile_quants[32][32]; - for (int cp = 0; cp < 16; cp++) { - for (int row = 0; row < 32; row++) { - uint8_t val = tile_src[cp * 32 + row]; - tile_quants[row][2 * cp + 0] = val & 0x0F; - tile_quants[row][2 * cp + 1] = val >> 4; - } - } + size_t slice_size = ne1 * ggml_row_size(t->type, ne0); + size_t row_size_bytes = ggml_row_size(t->type, ne0); + int64_t start_slice = offset / slice_size; + int64_t end_slice = (offset + size + slice_size - 1) / slice_size; + if (end_slice > ne2 * ne3) { + end_slice = ne2 * ne3; + } + + for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { + size_t cur_start_byte = (std::max)(offset, (size_t) slice_idx * slice_size); + size_t cur_end_byte = (std::min)(offset + size, (size_t) (slice_idx + 1) * slice_size); + size_t slice_offset_start = cur_start_byte - (size_t) slice_idx * slice_size; + size_t slice_offset_end = cur_end_byte - (size_t) slice_idx * slice_size; + + int64_t start_row = slice_offset_start / row_size_bytes; + int64_t end_row = (slice_offset_end + row_size_bytes - 1) / row_size_bytes; + end_row = (std::min)(end_row, ne1); + + int start_ct = start_row / 32; + int end_ct = (end_row + 31) / 32; + end_ct = (std::min)(end_ct, n_col_tiles); + block_mxfp4 * dst_slice = dst_matrix + (cur_start_byte - offset) / sizeof(block_mxfp4); + const uint8_t * matrix_src = (const uint8_t *) t->data + slice_idx * matrix_size; + + for (int ct = start_ct; ct < end_ct; ct++) { + for (int kt = 0; kt < n_k_tiles; kt++) { + const uint8_t * tile_src = matrix_src + (ct * n_k_tiles + kt) * tile_size; + + uint8_t tile_quants[32][32]; + for (int cp = 0; cp < 16; cp++) { for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - pack_mxfp4_quants(&dst_expert[r * (ne0 / 32) + kt], tile_quants[row], 0); - } + uint8_t val = tile_src[cp * 32 + row]; + tile_quants[row][2 * cp + 0] = val & 0x0F; + tile_quants[row][2 * cp + 1] = val >> 4; } + } - const uint8_t * scale_src = tile_src + 512; - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - dst_expert[r * (ne0 / 32) + kt].e = scale_src[row]; - } + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r >= start_row && r < end_row && kt < ne0 / 32) { + pack_mxfp4_quants(&dst_slice[(r - start_row) * (ne0 / 32) + kt], tile_quants[row], 0); + } + } + + const uint8_t * scale_src = tile_src + 512; + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r >= start_row && r < end_row && kt < ne0 / 32) { + dst_slice[(r - start_row) * (ne0 / 32) + kt].e = scale_src[row]; } } } } } - - GGML_UNUSED(size); } -static void ggml_backend_hexagon_buffer_set_tensor(ggml_backend_buffer_t buffer, - ggml_tensor * tensor, - const void * data, - size_t offset, - size_t size) { - auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context; - auto sess = sbuf->sess; - - HEX_VERBOSE("ggml-hex: %s set-tensor %s : data %p offset %zu size %zu\n", sess->c_name(), tensor->name, data, offset, size); - +static void repack_tensor_tiled(ggml_tensor * tensor, const void * data, size_t size) { switch (tensor->type) { case GGML_TYPE_Q4_0: - GGML_ASSERT(offset == 0); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_q4_0_tiled(tensor, data, size); + repack_q4_0_tiled(tensor, data, 0, size); break; case GGML_TYPE_Q4_1: - GGML_ASSERT(offset == 0); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_q4_1_tiled(tensor, data, size); + repack_q4_1_tiled(tensor, data, 0, size); break; case GGML_TYPE_Q8_0: - GGML_ASSERT(offset == 0); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_q8_0_tiled(tensor, data, size); + repack_q8_0_tiled(tensor, data, 0, size); break; case GGML_TYPE_IQ4_NL: - GGML_ASSERT(offset == 0); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - // IQ4_NL has identical block layout to Q4_0 (ggml_half d + uint8_t qs[16]) - repack_q4_0_tiled(tensor, data, size); + repack_q4_0_tiled(tensor, data, 0, size); break; case GGML_TYPE_MXFP4: - GGML_ASSERT(offset == 0); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_mxfp4_tiled(tensor, data, size); + repack_mxfp4_tiled(tensor, data, 0, size); break; default: - memcpy((char *) tensor->data + offset, data, size); break; } } +static void ggml_backend_hexagon_buffer_set_tensor(ggml_backend_buffer_t buffer, + ggml_tensor * tensor, + const void * data, + size_t offset, + size_t size) { + auto extra = (ggml_hexagon_tensor_extra *) tensor->extra; + auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context; + auto sess = sbuf->sess; + + if (ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { + extra->flags |= GGML_HEXAGON_TENSOR_WEIGHT; + if (ggml_hexagon_is_repack_type(tensor->type)) { + extra->flags |= GGML_HEXAGON_TENSOR_REPACK; + } + } + + HEX_VERBOSE("ggml-hex: %s set-tensor %s : data %p offset %zu size %zu usage %d flags 0x%x\n", + sess->c_name(), tensor->name, data, offset, size, (int) buffer->usage, extra->flags); + + if ((extra->flags & GGML_HEXAGON_TENSOR_REPACK) == 0) { + memcpy((char *) tensor->data + offset, data, size); + return; + } + + if (offset == 0 && size == ggml_nbytes(tensor) && extra->shadow_buf.empty()) { + repack_tensor_tiled(tensor, data, size); + return; + } + + if (extra->shadow_buf.size() < ggml_nbytes(tensor)) { + extra->shadow_buf.resize(ggml_nbytes(tensor)); + } + memcpy(extra->shadow_buf.data() + offset, data, size); + extra->shadow_size += size; + + if (extra->shadow_size >= ggml_nbytes(tensor)) { + repack_tensor_tiled(tensor, extra->shadow_buf.data(), extra->shadow_buf.size()); + extra->shadow_buf.clear(); + extra->shadow_buf.shrink_to_fit(); + extra->shadow_size = 0; + } +} + static void ggml_backend_hexagon_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { - auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context; - auto sess = sbuf->sess; + auto extra = (ggml_hexagon_tensor_extra *) tensor->extra; + auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context; + auto sess = sbuf->sess; - HEX_VERBOSE("ggml-hex: %s get-tensor %s : data %p offset %zu size %zu\n", sess->c_name(), tensor->name, data, offset, size); + HEX_VERBOSE("ggml-hex: %s get-tensor %s : data %p offset %zu size %zu usage %d flags 0x%x\n", + sess->c_name(), tensor->name, data, offset, size, (int) buffer->usage, extra->flags); + + if ((extra->flags & GGML_HEXAGON_TENSOR_REPACK) == 0) { + memcpy(data, (const char *) tensor->data + offset, size); + return; + } switch (tensor->type) { case GGML_TYPE_Q4_0: GGML_ASSERT(offset == 0); GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_tiled_q4_0(data, tensor, size); + repack_tiled_q4_0(data, tensor, offset, size); break; case GGML_TYPE_Q4_1: GGML_ASSERT(offset == 0); GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_tiled_q4_1(data, tensor, size); + repack_tiled_q4_1(data, tensor, offset, size); break; case GGML_TYPE_Q8_0: GGML_ASSERT(offset == 0); GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_tiled_q8_0(data, tensor, size); + repack_tiled_q8_0(data, tensor, offset, size); break; case GGML_TYPE_IQ4_NL: GGML_ASSERT(offset == 0); GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_tiled_q4_0(data, tensor, size); + repack_tiled_q4_0(data, tensor, offset, size); break; case GGML_TYPE_MXFP4: GGML_ASSERT(offset == 0); GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_tiled_mxfp4(data, tensor, size); + repack_tiled_mxfp4(data, tensor, offset, size); break; default: @@ -1035,11 +1321,121 @@ static bool ggml_backend_hexagon_buffer_cpy_tensor(ggml_backend_buffer_t bu GGML_UNUSED(dst); } +static void ggml_backend_hexagon_buffer_set_tensor_2d(ggml_backend_buffer_t buffer, + ggml_tensor * tensor, + const void * data, + size_t offset, + size_t size, + size_t n_copies, + size_t stride_tensor, + size_t stride_data) { + auto extra = (ggml_hexagon_tensor_extra *) tensor->extra; + auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context; + auto sess = sbuf->sess; + + if (ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { + extra->flags |= GGML_HEXAGON_TENSOR_WEIGHT; + if (ggml_hexagon_is_repack_type(tensor->type)) { + extra->flags |= GGML_HEXAGON_TENSOR_REPACK; + } + } + + HEX_VERBOSE("ggml-hex: %s set-tensor-2d %s : data %p offset %zu size %zu n_copies %zu stride_tensor %zu stride_data %zu usage %d flags 0x%x\n", + sess->c_name(), tensor->name, data, offset, size, n_copies, stride_tensor, stride_data, (int) buffer->usage, extra->flags); + + if ((extra->flags & GGML_HEXAGON_TENSOR_REPACK) == 0) { + for (size_t i = 0; i < n_copies; i++) { + memcpy((uint8_t *) tensor->data + offset + i * stride_tensor, (const uint8_t *) data + i * stride_data, size); + } + return; + } + + if (extra->shadow_buf.size() < ggml_nbytes(tensor)) { + extra->shadow_buf.resize(ggml_nbytes(tensor)); + } + for (size_t i = 0; i < n_copies; i++) { + memcpy(extra->shadow_buf.data() + offset + i * stride_tensor, (const uint8_t *) data + i * stride_data, size); + } + extra->shadow_size += n_copies * size; + + if (extra->shadow_size >= ggml_nbytes(tensor)) { + repack_tensor_tiled(tensor, extra->shadow_buf.data(), extra->shadow_buf.size()); + extra->shadow_buf.clear(); + extra->shadow_buf.shrink_to_fit(); + extra->shadow_size = 0; + } +} + +static void ggml_backend_hexagon_buffer_get_tensor_2d(ggml_backend_buffer_t buffer, + const ggml_tensor * tensor, + void * data, + size_t offset, + size_t size, + size_t n_copies, + size_t stride_tensor, + size_t stride_data) { + auto extra = (ggml_hexagon_tensor_extra *) tensor->extra; + auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context; + auto sess = sbuf->sess; + + HEX_VERBOSE("ggml-hex: %s get-tensor-2d %s : data %p offset %zu size %zu n_copies %zu stride_tensor %zu stride_data %zu usage %d\n", + sess->c_name(), tensor->name, data, offset, size, n_copies, stride_tensor, stride_data, (int) buffer->usage); + + if ((extra->flags & GGML_HEXAGON_TENSOR_REPACK) == 0) { + for (size_t i = 0; i < n_copies; i++) { + memcpy((uint8_t *)data + i * stride_data, (const uint8_t *)tensor->data + offset + i * stride_tensor, size); + } + return; + } + + size_t temp_size = n_copies > 0 ? (n_copies - 1) * stride_tensor + size : 0; + size_t slice_size = tensor->ne[1] * ggml_row_size(tensor->type, tensor->ne[0]); + size_t slice_offset = offset % slice_size; + size_t row_size_bytes = ggml_row_size(tensor->type, tensor->ne[0]); + + GGML_ASSERT((slice_offset % row_size_bytes) == 0 && "offset must be aligned to row boundary"); + GGML_ASSERT((temp_size % row_size_bytes) == 0 && "temp_size must be a multiple of row size"); + GGML_ASSERT((slice_offset / row_size_bytes) % 32 == 0 && "offset must be aligned to tile size (32 rows)"); + GGML_ASSERT((offset + temp_size) <= ggml_nbytes(tensor)); + + std::vector temp_buf(temp_size); + + switch (tensor->type) { + case GGML_TYPE_Q4_0: + repack_tiled_q4_0(temp_buf.data(), tensor, offset, temp_size); + break; + + case GGML_TYPE_Q4_1: + repack_tiled_q4_1(temp_buf.data(), tensor, offset, temp_size); + break; + + case GGML_TYPE_Q8_0: + repack_tiled_q8_0(temp_buf.data(), tensor, offset, temp_size); + break; + + case GGML_TYPE_IQ4_NL: + repack_tiled_q4_0(temp_buf.data(), tensor, offset, temp_size); + break; + + case GGML_TYPE_MXFP4: + repack_tiled_mxfp4(temp_buf.data(), tensor, offset, temp_size); + break; + + default: + memcpy(temp_buf.data(), (const uint8_t *) tensor->data + offset, temp_size); + break; + } + + for (size_t i = 0; i < n_copies; i++) { + memcpy((uint8_t *) data + i * stride_data, temp_buf.data() + i * stride_tensor, size); + } +} + static void ggml_backend_hexagon_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context; auto sess = sbuf->sess; - HEX_VERBOSE("ggml-hex: %s clear-buff base %p size %zu\n", sess->c_name(), (void *) sbuf->base, sbuf->size); - memset(sbuf->base, value, sbuf->size); + HEX_VERBOSE("ggml-hex: %s clear-buff base %p size %zu\n", sess->c_name(), (void *) sbuf->base(), sbuf->size()); + memset(sbuf->base(), value, sbuf->size()); } static ggml_backend_buffer_i ggml_backend_hexagon_buffer_interface = { @@ -1049,6 +1445,40 @@ static ggml_backend_buffer_i ggml_backend_hexagon_buffer_interface = { /* .memset_tensor = */ NULL, /* .set_tensor = */ ggml_backend_hexagon_buffer_set_tensor, /* .get_tensor = */ ggml_backend_hexagon_buffer_get_tensor, + /* .set_tensor_2d = */ ggml_backend_hexagon_buffer_set_tensor_2d, + /* .get_tensor_2d = */ ggml_backend_hexagon_buffer_get_tensor_2d, + /* .cpy_tensor = */ ggml_backend_hexagon_buffer_cpy_tensor, + /* .clear = */ ggml_backend_hexagon_buffer_clear, + /* .reset = */ NULL, +}; + +// ** backend buffer type + +static void ggml_backend_hexagon_host_buffer_set_tensor(ggml_backend_buffer_t buffer, + ggml_tensor * tensor, + const void * data, + size_t offset, + size_t size) { + memcpy((char *) tensor->data + offset, data, size); + GGML_UNUSED(buffer); +} + +static void ggml_backend_hexagon_host_buffer_get_tensor(ggml_backend_buffer_t buffer, + const ggml_tensor * tensor, + void * data, + size_t offset, + size_t size) { + memcpy(data, (const char *) tensor->data + offset, size); + GGML_UNUSED(buffer); +} + +static ggml_backend_buffer_i ggml_backend_hexagon_host_buffer_interface = { + /* .free_buffer = */ ggml_backend_hexagon_buffer_free_buffer, + /* .get_base = */ ggml_backend_hexagon_buffer_get_base, + /* .init_tensor = */ ggml_backend_hexagon_buffer_init_tensor, + /* .memset_tensor = */ NULL, + /* .set_tensor = */ ggml_backend_hexagon_host_buffer_set_tensor, + /* .get_tensor = */ ggml_backend_hexagon_host_buffer_get_tensor, /* .set_tensor_2d = */ NULL, /* .get_tensor_2d = */ NULL, /* .cpy_tensor = */ ggml_backend_hexagon_buffer_cpy_tensor, @@ -1066,24 +1496,22 @@ static ggml_backend_buffer_t ggml_backend_hexagon_buffer_type_alloc_buffer( ggml_backend_buffer_type_t buffer_type, size_t size) { auto sess = static_cast(buffer_type->context)->sess; try { - size += 4 * 1024; // guard page - ggml_hexagon_shared_buffer * sbuf = new ggml_hexagon_shared_buffer(sess, size); + ggml_hexagon_shared_buffer * sbuf = new ggml_hexagon_shared_buffer(sess, size, false, GGML_HEXAGON_FENCE_BUFFER_SIZE); return ggml_backend_buffer_init(buffer_type, ggml_backend_hexagon_buffer_interface, sbuf, size); } catch (const std::exception & exc) { - GGML_LOG_ERROR("ggml-hex: %s failed to allocate buffer context (host): %s\n", sess->c_name(), exc.what()); + GGML_LOG_ERROR("ggml-hex: %s failed to allocate device buffer context: %s\n", sess->c_name(), exc.what()); return nullptr; } } -static ggml_backend_buffer_t ggml_backend_hexagon_repack_buffer_type_alloc_buffer( +static ggml_backend_buffer_t ggml_backend_hexagon_host_buffer_type_alloc_buffer( ggml_backend_buffer_type_t buffer_type, size_t size) { auto sess = static_cast(buffer_type->context)->sess; try { - size += 4 * 1024; // guard page - ggml_hexagon_shared_buffer * sbuf = new ggml_hexagon_shared_buffer(sess, size); - return ggml_backend_buffer_init(buffer_type, ggml_backend_hexagon_buffer_interface, sbuf, size); + ggml_hexagon_shared_buffer * sbuf = new ggml_hexagon_shared_buffer(sess, size, false, GGML_HEXAGON_FENCE_BUFFER_SIZE); + return ggml_backend_buffer_init(buffer_type, ggml_backend_hexagon_host_buffer_interface, sbuf, size); } catch (const std::exception & exc) { - GGML_LOG_ERROR("ggml-hex: %s failed to allocate buffer context (repack): %s\n", sess->c_name(), exc.what()); + GGML_LOG_ERROR("ggml-hex: %s failed to allocate host buffer context: %s\n", sess->c_name(), exc.what()); return nullptr; } } @@ -1094,7 +1522,7 @@ static size_t ggml_backend_hexagon_buffer_type_get_alignment(ggml_backend_buffer } static size_t ggml_backend_hexagon_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const struct ggml_tensor * t) { - if (t->type == GGML_TYPE_Q4_0 || t->type == GGML_TYPE_Q4_1 || t->type == GGML_TYPE_Q8_0 || t->type == GGML_TYPE_IQ4_NL || t->type == GGML_TYPE_MXFP4) { + if (ggml_hexagon_is_repack_type(t->type)) { int64_t ne0 = hex_round_up(t->ne[0], 32); int64_t ne1 = hex_round_up(t->ne[1], 32); int64_t ne2 = t->ne[2]; @@ -1112,14 +1540,12 @@ static size_t ggml_backend_hexagon_buffer_type_get_max_size(ggml_backend_buffer_ } static bool ggml_backend_hexagon_buffer_type_is_host(ggml_backend_buffer_type_t buft) { - return opt_hostbuf; - + return false; GGML_UNUSED(buft); } -static bool ggml_backend_hexagon_repack_buffer_type_is_host(ggml_backend_buffer_type_t buft) { - return false; - +static bool ggml_backend_hexagon_host_buffer_type_is_host(ggml_backend_buffer_type_t buft) { + return true; GGML_UNUSED(buft); } @@ -1132,26 +1558,19 @@ static ggml_backend_buffer_type_i ggml_backend_hexagon_buffer_type_interface = { /* .is_host = */ ggml_backend_hexagon_buffer_type_is_host, }; -static ggml_backend_buffer_type_i ggml_backend_hexagon_repack_buffer_type_interface = { +static ggml_backend_buffer_type_i ggml_backend_hexagon_host_buffer_type_interface = { /* .get_name = */ ggml_backend_hexagon_buffer_type_name, - /* .alloc_buffer = */ ggml_backend_hexagon_repack_buffer_type_alloc_buffer, + /* .alloc_buffer = */ ggml_backend_hexagon_host_buffer_type_alloc_buffer, /* .get_alignment = */ ggml_backend_hexagon_buffer_type_get_alignment, /* .get_max_size = */ ggml_backend_hexagon_buffer_type_get_max_size, /* .get_alloc_size = */ ggml_backend_hexagon_buffer_type_get_alloc_size, - /* .is_host = */ ggml_backend_hexagon_repack_buffer_type_is_host, + /* .is_host = */ ggml_backend_hexagon_host_buffer_type_is_host, }; static bool ggml_backend_buffer_is_hexagon(const struct ggml_backend_buffer * b) { return b->buft->iface.get_alignment == ggml_backend_hexagon_buffer_type_get_alignment; } -static inline bool ggml_backend_buffer_is_hexagon_repack(const struct ggml_backend_buffer * b) { - if (!opt_hostbuf) { - return ggml_backend_buffer_is_hexagon(b); - } - return b->buft->iface.alloc_buffer == ggml_backend_hexagon_repack_buffer_type_alloc_buffer; -} - struct ggml_hexagon_opbatch { ggml_hexagon_session* sess; @@ -1165,8 +1584,6 @@ struct ggml_hexagon_opbatch { std::unordered_map t_map; // tensor ptr to index std::unordered_multimap d_map; // tensor data to index - - unsigned int n_bufs; // num buffers in the batch unsigned int n_tens; // num tensors ... unsigned int n_ops; // num ops ... @@ -1186,6 +1603,7 @@ struct ggml_hexagon_opbatch { b_map.clear(); t_map.clear(); d_map.clear(); + ops.resize(n_ops_max); } ggml_hexagon_opbatch(ggml_hexagon_session *sess, size_t batch_size, size_t max_vmem) { @@ -1218,39 +1636,39 @@ struct ggml_hexagon_opbatch { // add buffer and return its index int add_buffer(ggml_hexagon_shared_buffer * sbuf) { // Lookup by fd - auto it = b_map.find(sbuf->fd); + auto it = b_map.find(sbuf->fd()); if (it != b_map.end()) { return it->second; } // Add new buffer to the batch int bi = n_bufs++; GGML_ASSERT(n_bufs < HTP_OP_MAX_BUFS); - b_map.insert({sbuf->fd, bi}); + b_map.insert({sbuf->fd(), bi}); htp_buf_desc &b = h_bufs[bi]; - b.base = (uint64_t) sbuf->base; - b.fd = sbuf->fd; - b.size = sbuf->size; + b.base = (uint64_t) sbuf->base(); + b.fd = sbuf->fd(); + b.size = sbuf->size(); b_vmem += b.size; - HEX_VERBOSE("ggml-hex: %s add-buffer #%u : fd %d base %p size %zu : vmem %zu\n", sess->c_name(), bi, b.fd, (void*) sbuf->base, (size_t) b.size, b_vmem); + HEX_VERBOSE("ggml-hex: %s add-buffer #%u : fd %d base %p size %zu : vmem %zu\n", sess->c_name(), bi, b.fd, (void*) sbuf->base(), (size_t) b.size, b_vmem); return bi; } - - bool same_shape(const htp_tensor * h, const ggml_tensor * t) const { + auto extra = (ggml_hexagon_tensor_extra *) t->extra; + int64_t ne0 = t->ne[0]; int64_t ne1 = t->ne[1]; - const bool is_repack = ggml_backend_buffer_is_hexagon_repack(t->buffer) && ggml_hexagon_is_repack_type(t->type); + const bool is_repack = (extra->flags & GGML_HEXAGON_TENSOR_REPACK) != 0; if (is_repack) { ne0 = hex_round_up(ne0, 32); ne1 = hex_round_up(ne1, 32); } int64_t nb1 = is_repack ? ggml_row_size(t->type, ne0) : t->nb[1]; - int64_t nb2 = is_repack ? nb1 * ne1 : t->nb[2]; + int64_t nb2 = is_repack ? nb1 * ne1 : t->nb[2]; int64_t nb3 = is_repack ? nb2 * t->ne[2] : t->nb[3]; return (h->type == t->type) && @@ -1260,7 +1678,8 @@ struct ggml_hexagon_opbatch { // add tensor and return its index int add_tensor(const ggml_tensor * t) { - auto sbuf = static_cast(t->buffer->context); + auto extra = (ggml_hexagon_tensor_extra *) t->extra; + auto sbuf = static_cast(t->buffer->context); // First lookup by tensor data auto range = d_map.equal_range(t->data); @@ -1280,7 +1699,7 @@ struct ggml_hexagon_opbatch { t_map.insert({t, ti}); d_map.insert({t->data, ti}); - uint64_t t_offset = (uint8_t *) t->data - sbuf->base; + uint64_t t_offset = (uint8_t *) t->data - sbuf->base(); size_t t_size = ggml_nbytes(t); htp_tensor &h = h_tens[ti]; @@ -1289,7 +1708,7 @@ struct ggml_hexagon_opbatch { h.data = t_offset; h.type = t->type; - const bool is_repack = ggml_backend_buffer_is_hexagon_repack(t->buffer) && ggml_hexagon_is_repack_type(t->type); + const bool is_repack = (extra->flags & GGML_HEXAGON_TENSOR_REPACK) != 0; if (is_repack) { h.ne[0] = hex_round_up(t->ne[0], 32); h.ne[1] = hex_round_up(t->ne[1], 32); @@ -1308,11 +1727,15 @@ struct ggml_hexagon_opbatch { h.nb[0] = t->nb[0]; h.nb[1] = t->nb[1]; h.nb[2] = t->nb[2]; h.nb[3] = t->nb[3]; } - - h.flags = 0; - if (ggml_backend_buffer_get_usage(t->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { - h.flags |= HTP_TENSOR_COMPUTE; + if ((extra->flags & GGML_HEXAGON_TENSOR_WEIGHT) != 0) { + h.flags |= HTP_TENSOR_WEIGHT; + } + if ((extra->flags & GGML_HEXAGON_TENSOR_REPACK) != 0) { + h.flags |= HTP_TENSOR_REPACK; + } + if ((extra->flags & GGML_HEXAGON_TENSOR_FENCE) != 0) { + h.flags |= HTP_TENSOR_FENCE; } HEX_VERBOSE("ggml-hex: %s add-tensor #%u %s : bi %d data %p offset %zu size %zu flags 0x%x : %zu:%zu:%zu:%zu\n", sess->c_name(), @@ -1336,8 +1759,8 @@ struct ggml_hexagon_opbatch { extra_tens++; auto sbuf = static_cast(t->buffer->context); - if (!b_map.count(sbuf->fd)) { - extra_vmem += sbuf->size; + if (!b_map.count(sbuf->fd())) { + extra_vmem += sbuf->size(); extra_bufs += 1; } } @@ -1372,32 +1795,497 @@ struct ggml_hexagon_opbatch { o.opcode = node.opcode; o.flags = 0; - if (!(opt_opstage & HTP_OPSTAGE_COMPUTE)) { - o.flags |= HTP_OPFLAGS_SKIP_COMPUTE; - } + ggml_hexagon_dump_op_exec(sess->c_name(), ops[n], o.flags); + + auto inputs = node.get_inputs(); + for (unsigned int i=0; i < HTP_OP_MAX_INPUTS; i++) { + o.src[i] = (i < inputs.size() && inputs[i]) ? add_tensor(inputs[i]) : 0xffff; + } + + auto outputs = node.get_outputs(); + for (unsigned int i=0; i < HTP_OP_MAX_OUTPUTS; i++) { + o.dst[i] = (i < outputs.size() && outputs[i]) ? add_tensor(outputs[i]) : 0xffff; + } + } + + bool try_fuse_allreduce_add(const htp_opnode & node) { + if (n_ops == 0 || opt_ar_select != 2) return false; + if (node.opcode != HTP_OP_ADD) return false; + + htp_opnode & last_node = ops[n_ops - 1]; + if (last_node.opcode != HTP_OP_ALLREDUCE) return false; + + auto * ar_kparams = (struct htp_allreduce_kernel_params *) last_node.kernel_params; + const uint32_t rank = (uint32_t) ar_kparams->rank; + const ggml_tensor * ar_local = (rank < last_node.inputs.size()) ? last_node.inputs[rank] : nullptr; + const ggml_tensor * add_src0 = node.src0(); + const ggml_tensor * add_src1 = node.src1(); + + if (!add_src0 || !add_src1 || !ar_local) return false; + if (!ggml_hexagon_tensor_is_fuseable(ar_local)) return false; + + const ggml_tensor * res_tensor = nullptr; + if (add_src0 == ar_local || add_src0->data == ar_local->data) { + res_tensor = add_src1; + } else if (add_src1 == ar_local || add_src1->data == ar_local->data) { + res_tensor = add_src0; + } else { + return false; + } + + if (!res_tensor || !res_tensor->data) return false; + + if (ar_local->type != res_tensor->type) return false; + + const bool is_same_shape = (ar_local->ne[0] == res_tensor->ne[0] && ar_local->ne[1] == res_tensor->ne[1] && + ar_local->ne[2] == res_tensor->ne[2] && ar_local->ne[3] == res_tensor->ne[3]); + const bool is_row_bcast = (ar_local->ne[0] == res_tensor->ne[0] && + res_tensor->ne[1] == 1 && res_tensor->ne[2] == 1 && res_tensor->ne[3] == 1); + + if (!is_same_shape && !is_row_bcast) return false; + + if (is_same_shape) { + if (ar_local->nb[1] != res_tensor->nb[1] || ar_local->nb[2] != res_tensor->nb[2] || + ar_local->nb[3] != res_tensor->nb[3]) { + return false; + } + if (ggml_is_contiguous(ar_local) != ggml_is_contiguous(res_tensor)) { + return false; + } + } + if (ggml_is_contiguous(ar_local) != ggml_is_contiguous(node.dst())) { + return false; + } + + struct htp_allreduce_kernel_params new_kparams; + if (!ggml_hexagon_precompute_allreduce_params( + sess, node.dst(), (uint32_t) ar_kparams->rank, (uint32_t) ar_kparams->n_ranks, true, is_row_bcast, &new_kparams + )) { + HEX_VERBOSE("ggml-hex: %s skip ALLREDUCE_ADD fusion: solver failed\n", sess->c_name()); + return false; + } + + size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0; + auto fit_t = [&](const ggml_tensor * t) { + if (!t) return; + if (!t_map.count(t)) { + extra_tens++; + auto sbuf = static_cast(t->buffer->context); + if (!b_map.count(sbuf->fd())) { + extra_vmem += sbuf->size(); + extra_bufs += 1; + } + } + }; + fit_t(res_tensor); + fit_t(node.dst()); + if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) { + return false; + } + + last_node.opcode = HTP_OP_ALLREDUCE_ADD; + last_node.name = "ALLREDUCE+ADD"; + last_node.inputs.push_back(res_tensor); + last_node.outputs.clear(); + last_node.outputs.push_back(node.dst()); + last_node.fused.push_back(node.node); + memcpy(last_node.kernel_params, &new_kparams, sizeof(new_kparams)); + + htp_op_desc & o = h_ops[n_ops - 1]; + o.opcode = HTP_OP_ALLREDUCE_ADD; + memcpy(o.kernel_params, &new_kparams, sizeof(new_kparams)); + + const uint32_t n_ranks = (uint32_t) ar_kparams->n_ranks; + o.src[2 * n_ranks] = add_tensor(res_tensor); + o.dst[0] = add_tensor(node.dst()); + for (uint32_t d = 1; d < HTP_OP_MAX_OUTPUTS; d++) { + o.dst[d] = 0xffff; + } + + HEX_VERBOSE("ggml-hex: %s fused ALLREDUCE+ADD (#%u)\n", sess->c_name(), n_ops - 1); + return true; + } + + bool try_fuse_rms_norm_mul(const htp_opnode & node) { + if (n_ops == 0) return false; + if (node.opcode != HTP_OP_MUL) return false; + + htp_opnode & last_node = ops[n_ops - 1]; + if (last_node.opcode != HTP_OP_RMS_NORM) return false; + + const ggml_tensor * mul_src0 = node.src0(); + const ggml_tensor * mul_src1 = node.src1(); + const ggml_tensor * rms_out = last_node.dst(); + + if (!mul_src0 || !mul_src1 || !rms_out) return false; + if (!ggml_hexagon_tensor_is_fuseable(rms_out)) return false; + + const ggml_tensor * weight = nullptr; + if (mul_src0 == rms_out || mul_src0->data == rms_out->data) { + weight = mul_src1; + } else if (mul_src1 == rms_out || mul_src1->data == rms_out->data) { + weight = mul_src0; + } else { + return false; + } + + if (!weight || !weight->data) return false; + + const ggml_tensor * src0 = last_node.src0(); + if (!src0 || !src0->data) return false; + + if (src0->ne[0] != weight->ne[0] || src0->ne[0] != node.dst()->ne[0]) { + return false; + } + + const bool is_row_bcast = (weight->ne[1] == 1 && weight->ne[2] == 1 && weight->ne[3] == 1); + const bool is_same_shape = (src0->ne[0] == weight->ne[0] && src0->ne[1] == weight->ne[1] && + src0->ne[2] == weight->ne[2] && src0->ne[3] == weight->ne[3]); + if (!is_row_bcast && !is_same_shape) return false; + + if (!ggml_are_same_shape(src0, node.dst())) { + return false; + } + if (ggml_is_contiguous(src0) != ggml_is_contiguous(node.dst())) { + return false; + } + + struct htp_unary_kernel_params new_kparams; + ggml_hexagon_precompute_unary_params( + sess, HTP_OP_RMS_NORM_MUL, src0, weight, node.dst(), &new_kparams + ); + + if ((size_t) new_kparams.vtcm_size > sess->vtcm_size) { + HEX_VERBOSE("ggml-hex: %s skip RMS_NORM_MUL fusion: VTCM needed (%d) > budget (%zu)\n", + sess->c_name(), new_kparams.vtcm_size, sess->vtcm_size); + return false; + } + + size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0; + auto fit_t = [&](const ggml_tensor * t) { + if (!t) return; + if (!t_map.count(t)) { + extra_tens++; + auto sbuf = static_cast(t->buffer->context); + if (!b_map.count(sbuf->fd())) { + extra_vmem += sbuf->size(); + extra_bufs += 1; + } + } + }; + fit_t(weight); + fit_t(node.dst()); + if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) { + return false; + } + + last_node.opcode = HTP_OP_RMS_NORM_MUL; + last_node.name = "RMS_NORM+MUL"; + last_node.inputs.clear(); + last_node.inputs.push_back(src0); + last_node.inputs.push_back(weight); + last_node.outputs.clear(); + last_node.outputs.push_back(node.dst()); + last_node.fused.push_back(node.node); + memcpy(last_node.kernel_params, &new_kparams, sizeof(new_kparams)); + + htp_op_desc & o = h_ops[n_ops - 1]; + o.opcode = HTP_OP_RMS_NORM_MUL; + memcpy(o.kernel_params, &new_kparams, sizeof(new_kparams)); + + o.src[0] = add_tensor(src0); + o.src[1] = add_tensor(weight); + for (uint32_t s = 2; s < HTP_OP_MAX_INPUTS; s++) { + o.src[s] = 0xffff; + } + o.dst[0] = add_tensor(node.dst()); + for (uint32_t d = 1; d < HTP_OP_MAX_OUTPUTS; d++) { + o.dst[d] = 0xffff; + } + + HEX_VERBOSE("ggml-hex: %s fused RMS_NORM+MUL (#%u)\n", sess->c_name(), n_ops - 1); + return true; + } + + bool try_fuse_mul_mat_add(const htp_opnode & node) { + if (n_ops == 0) return false; + if (node.opcode != HTP_OP_ADD) return false; + + htp_opnode & last_node = ops[n_ops - 1]; + if (last_node.opcode != HTP_OP_MUL_MAT) return false; + + const ggml_tensor * add_src0 = node.src0(); + const ggml_tensor * add_src1 = node.src1(); + const ggml_tensor * mm_out = last_node.dst(); + + if (!add_src0 || !add_src1 || !mm_out) return false; + if (!ggml_hexagon_tensor_is_fuseable(mm_out)) return false; + + const ggml_tensor * src2 = nullptr; + if (add_src0 == mm_out || add_src0->data == mm_out->data) { + src2 = add_src1; + } else if (add_src1 == mm_out || add_src1->data == mm_out->data) { + src2 = add_src0; + } else { + return false; + } + + if (!src2 || !src2->data) return false; + + const ggml_tensor * src0 = last_node.src0(); + const ggml_tensor * src1 = last_node.src1(); + if (!src0 || !src1) return false; + + struct htp_mm_kernel_params kparams; + ggml_hexagon_precompute_fused_matmul_add_params(sess, src0, src1, src2, node.dst(), &kparams); + const int src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3]; + const bool can_fuse = (kparams.n_hmx > 0) || (src1_nrows == 1); + if (!can_fuse) return false; + + if ((size_t) kparams.vtcm_size > sess->vtcm_size) { + HEX_VERBOSE("ggml-hex: %s skip MUL_MAT_ADD fusion: VTCM needed (%d) > budget (%zu)\n", + sess->c_name(), kparams.vtcm_size, sess->vtcm_size); + return false; + } + + size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0; + auto fit_t = [&](const ggml_tensor * t) { + if (!t) return; + if (!t_map.count(t)) { + extra_tens++; + auto sbuf = static_cast(t->buffer->context); + if (!b_map.count(sbuf->fd())) { + extra_vmem += sbuf->size(); + extra_bufs += 1; + } + } + }; + fit_t(src2); + fit_t(node.dst()); + if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) { + return false; + } + + last_node.opcode = HTP_OP_MUL_MAT_ADD; + last_node.name = "MUL_MAT+ADD"; + last_node.inputs.clear(); + last_node.inputs.push_back(src0); + last_node.inputs.push_back(src1); + last_node.inputs.push_back(src2); + last_node.outputs.clear(); + last_node.outputs.push_back(node.dst()); + last_node.fused.push_back(node.node); + memcpy(last_node.kernel_params, &kparams, sizeof(kparams)); + + htp_op_desc & o = h_ops[n_ops - 1]; + o.opcode = HTP_OP_MUL_MAT_ADD; + memcpy(o.kernel_params, &kparams, sizeof(kparams)); + + o.src[0] = add_tensor(src0); + o.src[1] = add_tensor(src1); + o.src[2] = add_tensor(src2); + for (uint32_t s = 3; s < HTP_OP_MAX_INPUTS; s++) { + o.src[s] = 0xffff; + } + o.dst[0] = add_tensor(node.dst()); + for (uint32_t d = 1; d < HTP_OP_MAX_OUTPUTS; d++) { + o.dst[d] = 0xffff; + } + + HEX_VERBOSE("ggml-hex: %s fused MUL_MAT+ADD (#%u)\n", sess->c_name(), n_ops - 1); + return true; + } + + bool try_fuse_mul_mat_nx(const htp_opnode & node) { + if (n_ops == 0 || node.opcode != HTP_OP_MUL_MAT) return false; + if (!is_mergeable_mul_mat(node.node)) return false; + + const ggml_tensor * w_in = node.src0(); + const ggml_tensor * x_in = node.src1(); + const ggml_tensor * d_in = node.dst(); + if (!w_in || !x_in || !d_in) return false; + + htp_opnode & last_node = ops[n_ops - 1]; + + // Case 1: last_node is already MUL_MAT_NX + if (last_node.opcode == HTP_OP_MUL_MAT_NX) { + const uint32_t curr_n = (uint32_t) last_node.outputs.size(); + if (curr_n >= HTP_OP_MAX_OUTPUTS || curr_n + 1 >= HTP_OP_MAX_INPUTS) { + return false; + } + + const ggml_tensor * w0 = last_node.inputs[0]; + const ggml_tensor * x = last_node.inputs[curr_n]; + + if (x_in != x || w_in->type != w0->type || w_in->ne[0] != w0->ne[0]) { + return false; + } + + struct htp_mm_kernel_params kparams; + ggml_hexagon_precompute_fused_mmnx_params(sess, w0, x, curr_n + 1, &kparams); + if ((size_t) kparams.vtcm_size > sess->vtcm_size) { + HEX_VERBOSE("ggml-hex: %s skip NX fusion: VTCM needed (%d) > budget (%zu)\n", + sess->c_name(), kparams.vtcm_size, sess->vtcm_size); + return false; + } + + size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0; + auto fit_t = [&](const ggml_tensor * t) { + if (!t) return; + if (!t_map.count(t)) { + extra_tens++; + auto sbuf = static_cast(t->buffer->context); + if (!b_map.count(sbuf->fd())) { + extra_vmem += sbuf->size(); + extra_bufs += 1; + } + } + }; + fit_t(w_in); + fit_t(d_in); + if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) { + return false; + } + + last_node.inputs[curr_n] = w_in; + last_node.inputs.push_back(x); + last_node.outputs.push_back(d_in); + last_node.fused.push_back(node.node); + memcpy(last_node.kernel_params, &kparams, sizeof(kparams)); + + htp_op_desc & o = h_ops[n_ops - 1]; + memcpy(o.kernel_params, &kparams, sizeof(kparams)); + + for (uint32_t s = 0; s <= curr_n + 1; s++) { + o.src[s] = add_tensor(last_node.inputs[s]); + } + for (uint32_t s = curr_n + 2; s < HTP_OP_MAX_INPUTS; s++) { + o.src[s] = 0xffff; + } + for (uint32_t d = 0; d <= curr_n; d++) { + o.dst[d] = add_tensor(last_node.outputs[d]); + } + for (uint32_t d = curr_n + 1; d < HTP_OP_MAX_OUTPUTS; d++) { + o.dst[d] = 0xffff; + } + + HEX_VERBOSE("ggml-hex: %s fused MUL_MAT_NX (N=%u, #%u)\n", sess->c_name(), curr_n + 1, n_ops - 1); + return true; + } + + // Case 2: last_node is single MUL_MAT + if (last_node.opcode == HTP_OP_MUL_MAT) { + if (!is_mergeable_mul_mat_pair(last_node.node, node.node)) { + return false; + } + + const ggml_tensor * w0 = last_node.src0(); + const ggml_tensor * x = last_node.src1(); + const ggml_tensor * w1 = node.src0(); + if (!w0 || !x || !w1) return false; - ggml_hexagon_dump_op_exec(sess->c_name(), ops[n], o.flags); + struct htp_mm_kernel_params kparams; + ggml_hexagon_precompute_fused_mmnx_params(sess, w0, x, 2, &kparams); + if ((size_t) kparams.vtcm_size > sess->vtcm_size) { + HEX_VERBOSE("ggml-hex: %s skip NX fusion: VTCM needed (%d) > budget (%zu)\n", + sess->c_name(), kparams.vtcm_size, sess->vtcm_size); + return false; + } - auto inputs = node.get_inputs(); - for (unsigned int i=0; i < HTP_OP_MAX_INPUTS; i++) { - o.src[i] = (i < inputs.size() && inputs[i]) ? add_tensor(inputs[i]) : 0xffff; - } + size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0; + auto fit_t = [&](const ggml_tensor * t) { + if (!t) return; + if (!t_map.count(t)) { + extra_tens++; + auto sbuf = static_cast(t->buffer->context); + if (!b_map.count(sbuf->fd())) { + extra_vmem += sbuf->size(); + extra_bufs += 1; + } + } + }; + fit_t(w1); + fit_t(node.dst()); + if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) { + return false; + } - auto outputs = node.get_outputs(); - for (unsigned int i=0; i < HTP_OP_MAX_OUTPUTS; i++) { - o.dst[i] = (i < outputs.size() && outputs[i]) ? add_tensor(outputs[i]) : 0xffff; + const ggml_tensor * dst_0 = last_node.dst(); + const ggml_tensor * dst_1 = node.dst(); + + last_node.opcode = HTP_OP_MUL_MAT_NX; + last_node.name = "MUL_MAT_NX"; + last_node.inputs.clear(); + last_node.inputs.push_back(w0); + last_node.inputs.push_back(w1); + last_node.inputs.push_back(x); + last_node.outputs.clear(); + last_node.outputs.push_back(dst_0); + last_node.outputs.push_back(dst_1); + last_node.fused.push_back(node.node); + memcpy(last_node.kernel_params, &kparams, sizeof(kparams)); + + htp_op_desc & o = h_ops[n_ops - 1]; + o.opcode = HTP_OP_MUL_MAT_NX; + memcpy(o.kernel_params, &kparams, sizeof(kparams)); + + o.src[0] = add_tensor(w0); + o.src[1] = add_tensor(w1); + o.src[2] = add_tensor(x); + for (uint32_t s = 3; s < HTP_OP_MAX_INPUTS; s++) { + o.src[s] = 0xffff; + } + o.dst[0] = add_tensor(dst_0); + o.dst[1] = add_tensor(dst_1); + for (uint32_t d = 2; d < HTP_OP_MAX_OUTPUTS; d++) { + o.dst[d] = 0xffff; + } + + HEX_VERBOSE("ggml-hex: %s fused MUL_MAT_NX (N=2, #%u)\n", sess->c_name(), n_ops - 1); + return true; } + + return false; } - void finalize_ranges() { +enum ggml_hexagon_fusion_flags { + GGML_HEXAGON_FUSE_ALLREDUCE_ADD = (1 << 1), // 2 + GGML_HEXAGON_FUSE_RMS_NORM_MUL = (1 << 2), // 4 + GGML_HEXAGON_FUSE_MUL_MAT_ADD = (1 << 3), // 8 + GGML_HEXAGON_FUSE_MUL_MAT_NX = (1 << 4), // 16 +}; + +static inline bool ggml_hexagon_is_fusion_enabled(int flag) { + if (opt_opfusion <= 0) return false; + if (opt_opfusion == 1) return true; // 1 enables all + return (opt_opfusion & flag) != 0; +} + + bool try_fuse(const htp_opnode & node) { + if (!opt_opfusion) return false; + if (ggml_hexagon_is_fusion_enabled(GGML_HEXAGON_FUSE_ALLREDUCE_ADD) && try_fuse_allreduce_add(node)) return true; + if (ggml_hexagon_is_fusion_enabled(GGML_HEXAGON_FUSE_RMS_NORM_MUL) && try_fuse_rms_norm_mul(node)) return true; + if (ggml_hexagon_is_fusion_enabled(GGML_HEXAGON_FUSE_MUL_MAT_ADD) && try_fuse_mul_mat_add(node)) return true; + if (ggml_hexagon_is_fusion_enabled(GGML_HEXAGON_FUSE_MUL_MAT_NX) && try_fuse_mul_mat_nx(node)) return true; + return false; } }; +struct ggml_hexagon_registry { + ggml_hexagon_registry(ggml_backend_reg_t reg); + ~ggml_hexagon_registry(); + + ggml_backend_device devices[GGML_HEXAGON_MAX_SESSIONS]; +}; + struct ggml_hexagon_opqueue { // Shared buffer for storing batches ggml_hexagon_shared_buffer *shm_buf; size_t shm_blk_size; + uint64_t req_seq = 0; + uint64_t rsp_seq = 0; + using opvec = std::vector; std::queue done; // completed batch ids @@ -1429,8 +2317,8 @@ struct ggml_hexagon_opqueue { for (unsigned int i = 0; i < depth; i++) { done.push(i); } if (opt_verbose) { - GGML_LOG_INFO("ggml-hex: %s allocated op-queue : batch-size %zu depth %zu shm-size %zu shm-block-size %zu\n", - sess->c_name(), batch_size, depth, shm_buf->size, shm_blk_size); + GGML_LOG_INFO("ggml-hex: %s allocated opqueue : batch-size %zu depth %zu shm-size %zu shm-block-size %zu\n", + sess->c_name(), batch_size, depth, shm_buf->size(), shm_blk_size); } } @@ -1453,8 +2341,9 @@ struct ggml_hexagon_opqueue { req.n_bufs = op_batch->n_bufs; req.n_tensors = op_batch->n_tens; req.n_ops = op_batch->n_ops; + req.seq = ++req_seq; - op_cache[req.id] = op_batch->ops; + op_cache[req.id] = std::move(op_batch->ops); start_usec[req.id] = ggml_time_us(); const size_t b_size = sizeof(htp_buf_desc) * req.n_bufs; @@ -1470,10 +2359,10 @@ struct ggml_hexagon_opqueue { req.n_traces = 0; } - dbuf.ptr = shm_buf->base + (req.id * shm_blk_size); - dbuf.fd = shm_buf->fd; + dbuf.ptr = shm_buf->base() + (req.id * shm_blk_size); + dbuf.fd = shm_buf->fd(); dbuf.flags = DSPQUEUE_BUFFER_FLAG_FLUSH_SENDER | DSPQUEUE_BUFFER_FLAG_INVALIDATE_RECIPIENT; - dbuf.offset = (uint8_t*) dbuf.ptr - (uint8_t*) shm_buf->base; + dbuf.offset = (uint8_t*) dbuf.ptr - (uint8_t*) shm_buf->base(); dbuf.size = b_size + t_size + o_size + p_size + tr_size; GGML_ASSERT(dbuf.size <= shm_blk_size); @@ -1487,7 +2376,7 @@ struct ggml_hexagon_opqueue { memcpy(t_ptr, (void *) op_batch->h_tens.data(), t_size); memcpy(o_ptr, (void *) op_batch->h_ops.data(), o_size); - HEX_VERBOSE("ggml-hex: %s op-queue push batch #%u : n-bufs %u n-tensors %u n-ops %u vmem %zu : b-size %zu t-size %zu o-size %zu m-size %zu\n", + HEX_VERBOSE("ggml-hex: %s opqueue-push batch #%u : n-bufs %u n-tensors %u n-ops %u vmem %zu : b-size %zu t-size %zu o-size %zu m-size %zu\n", shm_buf->sess->c_name(), req.id, req.n_bufs, req.n_tensors, req.n_ops, op_batch->b_vmem, b_size, t_size, o_size, (size_t) dbuf.size); @@ -1530,33 +2419,40 @@ struct ggml_hexagon_opqueue { const size_t m_size = b_size + t_size + o_size + p_size + tr_size; GGML_ASSERT(m_size <= shm_blk_size); - HEX_VERBOSE("ggml-hex: %s op-queue pop batch #%u : n-bufs %u n-tensors %u n-ops %u : m-size %zu b-size %zu t-size %zu o-size %zu\n", + HEX_VERBOSE("ggml-hex: %s opqueue-pop batch #%u : n-bufs %u n-tensors %u n-ops %u : m-size %zu b-size %zu t-size %zu o-size %zu\n", shm_buf->sess->c_name(), rsp.id, rsp.n_bufs, rsp.n_tensors, rsp.n_ops, (size_t) dbuf.size, b_size, t_size, o_size); uint8_t * m_ptr = (uint8_t*) dbuf.ptr; uint8_t * p_ptr = m_ptr + (b_size + t_size + o_size); - if (opt_profile && rsp.n_ops > 0) { + if (rsp.n_ops > 0) { auto & ops = op_cache[rsp.id]; - GGML_ASSERT(rsp.n_ops <= ops.size()); const htp_prof_desc * pd = (const htp_prof_desc *) p_ptr; - const htp_trace_desc * trace_events = nullptr; - if (opt_profile == 3) { trace_events = (const htp_trace_desc *) (p_ptr + p_size); } - ggml_hexagon_dump_batch_prof(shm_buf->sess->name, rsp); + if (opt_profile) { + ggml_hexagon_dump_batch_prof(shm_buf->sess->name, rsp); + } for (uint32_t i = 0; i < rsp.n_ops; i++) { - ggml_hexagon_dump_op_prof(shm_buf->sess->name, ops[i], pd[i]); + if (opt_profile) { + ggml_hexagon_dump_op_prof(shm_buf->sess->name, ops[i], pd[i]); + } } - ggml_hexagon_dump_trace_events(shm_buf->sess->name, rsp, trace_events, n_traces); + if (opt_profile) { + ggml_hexagon_dump_trace_events(shm_buf->sess->name, rsp, trace_events, n_traces); + } + } + + if (rsp.seq > rsp_seq) { + rsp_seq = rsp.seq; } } }; @@ -1601,10 +2497,8 @@ void ggml_hexagon_session::flush_pending(bool all) { } } -void ggml_hexagon_session::flush_batch() { - if (op_batch->empty()) { return; } - - op_batch->finalize_ranges(); +void ggml_hexagon_session::flush_batch(size_t min_ops) { + if (op_batch->n_ops < min_ops) { return; } htp_opbatch_req req {}; dspqueue_buffer dbuf{}; @@ -1625,17 +2519,267 @@ void ggml_hexagon_session::flush_batch() { } } +void ggml_hexagon_session::flush(bool all) { + flush_sync_peers(); + flush_batch(); + flush_pending(all); +} + void ggml_hexagon_session::enqueue_op(const htp_opnode & node) { + for (auto t : node.get_inputs()) { + if (t && t->buffer && ggml_backend_buffer_is_hexagon(t->buffer)) { + if (ggml_backend_hexagon_buffer_get_sess(t->buffer) != this) { + this->clone_buffer(static_cast(t->buffer->context)); + } + } + } + for (auto t : node.get_outputs()) { + if (t && t->buffer && ggml_backend_buffer_is_hexagon(t->buffer)) { + if (ggml_backend_hexagon_buffer_get_sess(t->buffer) != this) { + this->clone_buffer(static_cast(t->buffer->context)); + } + } + } + + if (opt_opfusion && op_batch->try_fuse(node)) { + return; + } + if (!op_batch->fit_op(node)) { flush_batch(); } op_batch->add_op(node); } -// Flush HTP response queue i.e wait for all outstanding requests to complete -void ggml_hexagon_session::flush(bool all) { +void ggml_hexagon_session::enqueue_cpy(const ggml_tensor * src, ggml_tensor * dst, const ggml_tensor * sync_tensor, uint32_t fence_seq) { + htp_opnode cpy_node(HTP_OP_CPY); + + ggml_tensor* node = cpy_node.add_dummy(*dst); + node->op = GGML_OP_CPY; + node->src[0] = const_cast(src); + node->src[1] = sync_tensor ? cpy_node.add_dummy(*sync_tensor) : nullptr; + if (sync_tensor) { + node->op_params[0] = (int32_t) fence_seq; + } + + cpy_node.init(node); + if (sync_tensor) { + cpy_node.name = "CPY+FENCE"; + } + this->enqueue_op(cpy_node); +} + +void ggml_hexagon_session::enqueue_fence(const ggml_tensor * sync_tensor, uint32_t fence_seq) { + htp_opnode sync_node(HTP_OP_FENCE); + + ggml_tensor* node = sync_node.add_dummy(*sync_tensor); + node->op = GGML_OP_NONE; + node->src[0] = node; + node->op_params[0] = (int32_t) fence_seq; + + sync_node.init(node); + sync_node.name = "FENCE"; + this->enqueue_op(sync_node); +} + +static bool ggml_hexagon_precompute_allreduce_params( + const struct ggml_hexagon_session * sess, + const struct ggml_tensor * dst, + uint32_t rank, + uint32_t n_ranks, + bool has_add, + bool is_row_bcast, + struct htp_allreduce_kernel_params * kparams +) { + memset(kparams, 0, sizeof(*kparams)); + kparams->rank = (int32_t) rank; + kparams->n_ranks = (int32_t) n_ranks; + kparams->is_row_bcast = (has_add && is_row_bcast) ? 1 : 0; + + const uint32_t n_bufs = n_ranks + 1 + (has_add ? 1 : 0); + const uint32_t nelem = (uint32_t) ggml_nelements(dst); + const uint32_t elem_size = (dst->type == GGML_TYPE_F16) ? sizeof(ggml_fp16_t) : sizeof(float); + const bool is_contiguous = ggml_is_contiguous(dst); + + const uint32_t ne0 = (uint32_t) dst->ne[0]; + const uint32_t ne1 = (uint32_t) (dst->ne[1] * dst->ne[2] * dst->ne[3]); + kparams->ne0 = (int32_t) ne0; + kparams->ne1 = (int32_t) ne1; + + const bool use_1d = is_contiguous && !(has_add && is_row_bcast && ne1 > 1); + + if (has_add) { + kparams->n_dsts = 1; + if (use_1d) { + kparams->rank_elem_start = 0; + kparams->rank_nelem = (int32_t) nelem; + } else { + kparams->rank_elem_start = 0; + kparams->rank_nelem = (int32_t) ne1; + } + } else { + kparams->n_dsts = (int32_t) n_ranks; + if (use_1d) { + const uint32_t rank_chunk_elems = hex_round_up((nelem + n_ranks - 1) / n_ranks, 128); + const uint32_t rank_elem_start = (std::min)(rank * rank_chunk_elems, nelem); + const uint32_t rank_elem_end = (std::min)(rank_elem_start + rank_chunk_elems, nelem); + const uint32_t rank_nelem = rank_elem_end - rank_elem_start; + kparams->rank_elem_start = (int32_t) rank_elem_start; + kparams->rank_nelem = (int32_t) rank_nelem; + } else { + const uint32_t rank_chunk_rows = (ne1 + n_ranks - 1) / n_ranks; + const uint32_t rank_r0 = (std::min)(rank * rank_chunk_rows, ne1); + const uint32_t rank_r1 = (std::min)(rank_r0 + rank_chunk_rows, ne1); + const uint32_t rank_nrows = rank_r1 - rank_r0; + kparams->rank_elem_start = (int32_t) rank_r0; + kparams->rank_nelem = (int32_t) rank_nrows; + } + } + + if (use_1d) { + const uint32_t rank_nelem = (uint32_t) kparams->rank_nelem; + const uint32_t n_threads = (std::min)((uint32_t) sess->n_threads, (std::max)(1u, rank_nelem / 128)); + kparams->n_threads = n_threads; + + uint32_t block_elems = 65536; + if (block_elems > rank_nelem / n_threads && rank_nelem / n_threads > 128) { + block_elems = hex_round_up(rank_nelem / (n_threads * 2), 128); + } + block_elems = (std::max)(128u, block_elems); + + kparams->block_elems = block_elems; + kparams->vtcm_size_per_thread = 2 * block_elems * elem_size; + kparams->vtcm_size = n_threads * n_bufs * kparams->vtcm_size_per_thread; + + while ((size_t) kparams->vtcm_size > sess->vtcm_size && block_elems > 128) { + const size_t max_bytes_per_buf = sess->vtcm_size / (n_threads * n_bufs * 2); + block_elems = (uint32_t) hex_align_down((size_t) (max_bytes_per_buf / elem_size), 128); + if (block_elems < 128) break; + kparams->block_elems = block_elems; + kparams->vtcm_size_per_thread = 2 * block_elems * elem_size; + kparams->vtcm_size = n_threads * n_bufs * kparams->vtcm_size_per_thread; + } + + if (sess->vtcm_size < (size_t) kparams->vtcm_size || block_elems < 128) { + HEX_VERBOSE("ggml-hex: %s allreduce 1D solver failed to fit VTCM (%d > %zu)\n", + sess->c_name(), kparams->vtcm_size, sess->vtcm_size); + return false; + } + + kparams->elems_per_thread = hex_round_up((rank_nelem + n_threads - 1) / n_threads, block_elems); + kparams->kernel_type = HTP_ALLREDUCE_KERNEL_DMA_1D; + return true; + } else { + const uint32_t rank_nrows = (uint32_t) kparams->rank_nelem; + const uint32_t n_threads = (std::min)((uint32_t) sess->n_threads, (std::max)(1u, rank_nrows)); + kparams->n_threads = n_threads; + + const uint32_t row_bytes = ne0 * elem_size; + const uint32_t row_size_aligned = (uint32_t) hex_align_up(row_bytes, 128); + kparams->row_size_aligned = row_size_aligned; + + const uint32_t nrows_per_thread = (rank_nrows + n_threads - 1) / n_threads; + uint32_t block_rows = (std::min)(128u, nrows_per_thread); + block_rows = (std::max)(1u, block_rows); + kparams->block_elems = block_rows; + + kparams->vtcm_size_per_thread = 2 * (block_rows * row_size_aligned); + kparams->vtcm_size = n_threads * n_bufs * kparams->vtcm_size_per_thread; + + while ((size_t) kparams->vtcm_size > sess->vtcm_size && block_rows > 1) { + const size_t max_rows_per_buf = sess->vtcm_size / (n_threads * n_bufs * 2 * row_size_aligned); + block_rows = (std::max)(1u, (uint32_t) max_rows_per_buf); + kparams->block_elems = block_rows; + kparams->vtcm_size_per_thread = 2 * (block_rows * row_size_aligned); + kparams->vtcm_size = n_threads * n_bufs * kparams->vtcm_size_per_thread; + if (max_rows_per_buf == 0) break; + } + + if (sess->vtcm_size < (size_t) kparams->vtcm_size || block_rows < 1) { + HEX_VERBOSE("ggml-hex: %s allreduce 2D solver failed to fit VTCM (%d > %zu)\n", + sess->c_name(), kparams->vtcm_size, sess->vtcm_size); + return false; + } + + kparams->elems_per_thread = nrows_per_thread; + kparams->kernel_type = HTP_ALLREDUCE_KERNEL_DMA_2D; + return true; + } +} + +void ggml_hexagon_session::enqueue_allreduce( + const ggml_tensor * dst, + const std::vector & src_tensors, + const std::vector & sync_tensors, + uint32_t rank, + uint32_t n_ranks, + uint32_t fence_seq_entry, + uint32_t fence_seq_exit +) { + htp_opnode ar_node(HTP_OP_ALLREDUCE); + + ggml_tensor* node = ar_node.add_dummy(*dst); + node->op = GGML_OP_NONE; + node->op_params[0] = (int32_t) fence_seq_entry; + node->op_params[1] = (int32_t) fence_seq_exit; + + ar_node.init(node); + + ar_node.inputs.clear(); + for (size_t i = 0; i < src_tensors.size(); i++) { + ar_node.inputs.push_back(src_tensors[i]); + } + for (size_t i = 0; i < sync_tensors.size(); i++) { + ar_node.inputs.push_back(ar_node.add_dummy(*sync_tensors[i])); + } + + ar_node.outputs.clear(); + for (size_t i = 0; i < src_tensors.size(); i++) { + ar_node.outputs.push_back(src_tensors[i]); + } + + ggml_hexagon_precompute_allreduce_params( + this, dst, rank, n_ranks, false, false, + (struct htp_allreduce_kernel_params *) ar_node.kernel_params + ); + + ar_node.name = "ALLREDUCE"; + this->enqueue_op(ar_node); +} + +void ggml_hexagon_session::wait_event(uint64_t seq) { + flush_sync_peers(); + HEX_VERBOSE("ggml-hex: %s opqueue-wait start: seq %llu, current rsp-seq %llu, pending %d\n", + this->name.c_str(), (unsigned long long)seq, (unsigned long long)op_queue->rsp_seq, (int)this->op_pending); + while (op_queue->rsp_seq < seq && this->op_pending > 0) { + this->flush_pending(false); + } + HEX_VERBOSE("ggml-hex: %s opqueue-wait end: seq %llu, current rsp-seq %llu, pending %d\n", + this->name.c_str(), (unsigned long long)seq, (unsigned long long)op_queue->rsp_seq, (int)this->op_pending); +} + +uint64_t ggml_hexagon_session::record_event() { flush_batch(); - flush_pending(all); + return op_queue->req_seq; +} + +bool ggml_hexagon_session::clone_buffer(const ggml_hexagon_shared_buffer *sbuf) +{ + if (this->cloned_buffers.find(sbuf->fd()) != this->cloned_buffers.end()) return true; + + HEX_VERBOSE("ggml-hex: %s clone-buffer: %s base %p size %zu fd %d\n", this->name.c_str(), + sbuf->c_name(), sbuf->base(), sbuf->size(), sbuf->fd()); + + auto clone = std::make_unique(this, *sbuf); + try { + clone->mmap(); + } catch (const std::exception & exc) { + GGML_LOG_ERROR("ggml-hex: %s lazy mapping of buffer context failed: %s\n", this->c_name(), exc.what()); + return false; + } + + this->cloned_buffers[sbuf->fd()] = std::move(clone); + return true; } static size_t ggml_hexagon_measure_max_vmem(ggml_hexagon_session *sess) { @@ -1668,37 +2812,44 @@ static size_t ggml_hexagon_measure_max_vmem(ggml_hexagon_session *sess) { } void ggml_hexagon_session::allocate(int dev_id) noexcept(false) { + const auto & config = opt_device_configs[dev_id]; + int phys_idx = config.physical_idx; + int virt_idx = config.virtual_idx; + this->valid_session = false; this->valid_handle = false; this->valid_queue = false; this->valid_iface = false; - this->domain_id = 3; // Default for CDSP, updated after the session is created - this->session_id = 0; // Default for CDSP, updated after the session is created + this->phys_idx = phys_idx; + this->virt_idx = virt_idx; + this->domain_id = get_domain_id(phys_idx); + this->session_id = 0; this->dev_id = dev_id; - this->name = std::string("HTP") + std::to_string(dev_id); - - this->op_pending = 0; + this->name = config.name; + this->op_pending = 0; GGML_LOG_DEBUG("ggml-hex: %s allocating new session\n", this->name.c_str()); domain * my_domain = htpdrv_get_domain(this->domain_id); if (my_domain == NULL) { - GGML_LOG_ERROR("ggml-hex: unable to get domain struct for CDSP\n"); + GGML_LOG_ERROR("ggml-hex: unable to get domain struct for CDSP (domain_id %d)\n", this->domain_id); throw std::runtime_error("ggml-hex: failed to get CDSP domain (see log for details)"); } - // Create new session - if (dev_id != 0) { + std::string dom_name = get_domain_name(phys_idx); + + // Create new session if virtual_idx > 0 + if (virt_idx > 0) { struct remote_rpc_reserve_new_session n; - n.domain_name_len = strlen(CDSP_DOMAIN_NAME); - n.domain_name = const_cast(CDSP_DOMAIN_NAME); + n.domain_name_len = dom_name.size(); + n.domain_name = const_cast(dom_name.c_str()); n.session_name = const_cast(this->name.c_str()); n.session_name_len = this->name.size(); int err = remote_session_control(FASTRPC_RESERVE_NEW_SESSION, (void *) &n, sizeof(n)); if (err != AEE_SUCCESS) { - GGML_LOG_ERROR("ggml-hex: failed to reserve new session %d : error 0x%x\n", dev_id, err); + GGML_LOG_ERROR("ggml-hex: failed to reserve new session %d (physical %d, virtual %d) : error 0x%x\n", dev_id, phys_idx, virt_idx, err); throw std::runtime_error("ggml-hex: remote_session_control(new-sess) failed (see log for details)"); } @@ -1717,8 +2868,8 @@ void ggml_hexagon_session::allocate(int dev_id) noexcept(false) { struct remote_rpc_get_uri u = {}; u.session_id = this->session_id; - u.domain_name = const_cast(CDSP_DOMAIN_NAME); - u.domain_name_len = strlen(CDSP_DOMAIN_NAME); + u.domain_name = const_cast(dom_name.c_str()); + u.domain_name_len = dom_name.size(); u.module_uri = const_cast(htp_uri); u.module_uri_len = strlen(htp_uri); u.uri = session_uri; @@ -1731,7 +2882,7 @@ void ggml_hexagon_session::allocate(int dev_id) noexcept(false) { snprintf(session_uri, htp_URI_domain_len, "%s%s", htp_uri, my_domain->uri); - GGML_LOG_WARN("ggml-hex: failed to get URI for session %d : error 0x%x. Falling back to single session URI: %s\n", dev_id, err, session_uri); + GGML_LOG_WARN("ggml-hex: failed to get URI for session %d (physical %d, virtual %d) : error 0x%x. Falling back to single session URI: %s\n", dev_id, phys_idx, virt_idx, err, session_uri); } } @@ -1899,14 +3050,17 @@ void ggml_hexagon_session::release() noexcept(true) { if (this->valid_handle) { htp_iface_close(this->handle); } + + this->cloned_buffers.clear(); } ggml_hexagon_session::ggml_hexagon_session(int dev_id, ggml_backend_dev_t dev) noexcept(false) { - buffer_type.device = dev; - repack_buffer_type.device = dev; + buffer_type.device = dev; + host_buffer_type.device = dev; op_batch = nullptr; op_queue = nullptr; + fence_seq = ((uintptr_t)this) & 0xFFFF; try { allocate(dev_id); @@ -1914,8 +3068,8 @@ ggml_hexagon_session::ggml_hexagon_session(int dev_id, ggml_backend_dev_t dev) n buffer_type.iface = ggml_backend_hexagon_buffer_type_interface; buffer_type.context = new ggml_backend_hexagon_buffer_type_context(this->name, this); - repack_buffer_type.iface = ggml_backend_hexagon_repack_buffer_type_interface; - repack_buffer_type.context = new ggml_backend_hexagon_buffer_type_context(this->name + "-REPACK", this); + host_buffer_type.iface = ggml_backend_hexagon_host_buffer_type_interface; + host_buffer_type.context = new ggml_backend_hexagon_buffer_type_context(this->name + "-HOST", this); } catch (const std::exception & exc) { release(); throw; @@ -1926,7 +3080,7 @@ ggml_hexagon_session::~ggml_hexagon_session() noexcept(true) { release(); delete static_cast(buffer_type.context); - delete static_cast(repack_buffer_type.context); + delete static_cast(host_buffer_type.context); } // ** backend interface @@ -1946,7 +3100,8 @@ static bool ggml_hexagon_flash_attn_is_hmx_eligible( return false; } - if (k->type != GGML_TYPE_F16 || v->type != GGML_TYPE_F16) { + if ((k->type != GGML_TYPE_F16 && k->type != GGML_TYPE_Q8_0) || + (v->type != GGML_TYPE_F16 && v->type != GGML_TYPE_Q8_0)) { return false; } @@ -2098,8 +3253,10 @@ static bool ggml_hexagon_supported_flash_attn_ext(const struct ggml_hexagon_sess const struct ggml_tensor * src4 = op->src[4]; const struct ggml_tensor * dst = op; - // Check for F16 support only as requested - if ((src0->type != GGML_TYPE_F16 && src0->type != GGML_TYPE_F32) || src1->type != GGML_TYPE_F16 || src2->type != GGML_TYPE_F16) { + // Check for F16/Q8_0 support + if ((src0->type != GGML_TYPE_F16 && src0->type != GGML_TYPE_F32) || + (src1->type != GGML_TYPE_F16 && src1->type != GGML_TYPE_Q8_0) || + (src2->type != GGML_TYPE_F16 && src2->type != GGML_TYPE_Q8_0)) { return false; } @@ -2352,7 +3509,7 @@ static void ggml_hexagon_precompute_hvx_mm_params( for (uint32_t d = max_prefetch; d >= 2; d /= 2) { htp_mm_hvx_vtcm_layout_build( &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0->nb[1], 0, src2_row_size, d, true, false, false + 0, src0->nb[1], 0, src2_row_size, d, true, false ); if (L.total_bytes <= vtcm_budget) { best_n_prefetch = d; @@ -2362,7 +3519,7 @@ static void ggml_hexagon_precompute_hvx_mm_params( if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) { htp_mm_hvx_vtcm_layout_build( &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0->nb[1], 0, src2_row_size, 2, true, false, false + 0, src0->nb[1], 0, src2_row_size, 2, true, false ); } kparams->n_prefetch = best_n_prefetch; @@ -2386,7 +3543,7 @@ static void ggml_hexagon_precompute_hvx_mm_params( for (uint32_t d = max_prefetch; d >= 2; d /= 2) { htp_mm_hvx_vtcm_layout_build( &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, d, false, false, false + dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, d, false, false ); if (L.total_bytes <= vtcm_budget) { best_n_prefetch = d; @@ -2396,7 +3553,7 @@ static void ggml_hexagon_precompute_hvx_mm_params( if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) { htp_mm_hvx_vtcm_layout_build( &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 2, false, false, false + dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 2, false, false ); } @@ -2420,7 +3577,7 @@ static void ggml_hexagon_precompute_hvx_mm_params( struct htp_mm_hvx_vtcm_layout L; htp_mm_hvx_vtcm_layout_build( &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false + dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false ); kparams->n_prefetch = 16; @@ -2440,7 +3597,7 @@ static void ggml_hexagon_precompute_hvx_mm_params( struct htp_mm_hvx_vtcm_layout L; htp_mm_hvx_vtcm_layout_build( &L, HTP_MM_KERNEL_HVX_F16_F16_VTCM, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false + dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false ); if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) { @@ -2460,7 +3617,7 @@ static void ggml_hexagon_precompute_hvx_mm_params( kparams->src1_row_size = src1->nb[1]; htp_mm_hvx_vtcm_layout_build( &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false + dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false ); kparams->vtcm_size = L.total_bytes; kparams->vtcm_src0_size = L.src0_bytes; @@ -2476,7 +3633,7 @@ static void ggml_hexagon_precompute_hvx_mm_params( struct htp_mm_hvx_vtcm_layout L; htp_mm_hvx_vtcm_layout_build( &L, HTP_MM_KERNEL_HVX_F32_F32_VTCM, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false + dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false ); if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) { @@ -2492,7 +3649,7 @@ static void ggml_hexagon_precompute_hvx_mm_params( kparams->src1_row_size = src1->nb[1]; htp_mm_hvx_vtcm_layout_build( &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false + dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false ); kparams->vtcm_size = L.total_bytes; kparams->vtcm_src0_size = L.src0_bytes; @@ -2642,80 +3799,99 @@ static void ggml_hexagon_precompute_unary_params( kparams->div_tpr = init_fastdiv_values(tiles_per_row); } -static void ggml_hexagon_precompute_fused_qkv_params( +static void ggml_hexagon_precompute_get_rows_params( const struct ggml_hexagon_session * sess, - const struct ggml_tensor * src0, // Wk - const struct ggml_tensor * src1, // x - struct htp_mm_kernel_params * kparams + const struct ggml_tensor * src0, + const struct ggml_tensor * src1, + const struct ggml_tensor * dst, + struct htp_get_rows_kernel_params * kparams ) { memset(kparams, 0, sizeof(*kparams)); - const int wtype = src0->type; - const bool is_repack = ggml_hexagon_is_repack_type((ggml_type) wtype); + const uint32_t ne00 = src0->ne[0]; + const uint32_t ne02 = src0->ne[2]; + const uint32_t ne03 = src0->ne[3]; - const int ne10 = src1->ne[0]; - const int src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3]; - const size_t src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); - const size_t src0_row_size = src0->nb[1]; + const uint32_t ne10 = src1->ne[0]; + const uint32_t ne11 = src1->ne[1]; + const uint32_t ne12 = src1->ne[2]; + const uint32_t nr = ne10 * ne11 * ne12; - uint32_t best_n_prefetch = 16; + const size_t nb01 = src0->nb[1]; + const size_t nb1 = dst->nb[1]; - if (is_repack) { - const uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16; - best_n_prefetch = 2; - for (uint32_t d = max_prefetch; d >= 2; d /= 2) { - struct htp_mm_hvx_vtcm_layout L; - htp_mm_hvx_vtcm_layout_build( - &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0_row_size, src1_row_size, 0, d, false, true, false - ); - if (L.total_bytes <= sess->vtcm_size) { - best_n_prefetch = d; - break; + const bool can_use_dma = (src0->type == dst->type) && (nb01 == nb1); + const bool use_dma = can_use_dma && (ne00 >= 2048); + + kparams->use_dma = use_dma ? 1 : 0; + + uint32_t chunks_per_row = 1; + uint32_t chunk_size = ne00; + uint32_t total_tasks = nr; + + if (use_dma) { + kparams->n_threads = (std::min)((uint32_t)sess->n_threads, nr); + kparams->tasks_per_thread = (nr + kparams->n_threads - 1) / kparams->n_threads; + } else { + if (src0->type == GGML_TYPE_F32 && nr < sess->n_threads) { + const uint32_t min_chunk_size = 1024; + uint32_t max_chunks = ne00 / min_chunk_size; + if (max_chunks == 0) { + max_chunks = 1; } + chunks_per_row = (std::min)((sess->n_threads + nr - 1) / nr, max_chunks); + chunk_size = (ne00 + chunks_per_row - 1) / chunks_per_row; + total_tasks = nr * chunks_per_row; } + kparams->n_threads = (std::min)(total_tasks, (uint32_t)sess->n_threads); + kparams->tasks_per_thread = (total_tasks + kparams->n_threads - 1) / kparams->n_threads; } - struct htp_mm_hvx_vtcm_layout L; - bool try_tiled = (opt_mm_select >= 2); + kparams->chunks_per_row = chunks_per_row; + kparams->chunk_size = chunk_size; + kparams->total_tasks = total_tasks; - // Test tiled first - htp_mm_hvx_vtcm_layout_build( - &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0_row_size, src1_row_size, 0, best_n_prefetch, false, true, false - ); + kparams->div_ne10 = init_fastdiv_values(ne10); + kparams->div_ne10_ne11 = init_fastdiv_values(ne10 * ne11); + kparams->div_chunks_per_row = init_fastdiv_values(chunks_per_row); + kparams->div_ne02 = init_fastdiv_values(ne02); + kparams->div_ne03 = init_fastdiv_values(ne03); - if (try_tiled && L.total_bytes <= sess->vtcm_size) { - kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW; - kparams->vtcm_src0_size = L.src0_bytes; - kparams->vtcm_src1_size = L.src1_bytes; - kparams->vtcm_src2_size = L.src2_bytes; - kparams->vtcm_src3_size = L.src3_bytes; - kparams->vtcm_dst_size = L.dst_bytes; - kparams->vtcm_size = L.total_bytes; - kparams->n_prefetch = best_n_prefetch; - } else { - kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT; - size_t flat_src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); + struct htp_get_rows_vtcm_layout vtcm_layout; + htp_get_rows_vtcm_layout_build(&vtcm_layout, src0->type, ne00, kparams->n_threads); + kparams->vtcm_size = vtcm_layout.total_bytes; +} - htp_mm_hvx_vtcm_layout_build( - &L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0_row_size, flat_src1_row_size, 0, best_n_prefetch, false, true, false - ); - kparams->vtcm_src0_size = L.src0_bytes; - kparams->vtcm_src1_size = L.src1_bytes; - kparams->vtcm_src2_size = L.src2_bytes; - kparams->vtcm_src3_size = L.src3_bytes; - kparams->vtcm_dst_size = L.dst_bytes; - kparams->vtcm_size = L.total_bytes; - kparams->n_prefetch = best_n_prefetch; - } +static void ggml_hexagon_precompute_set_rows_params( + const struct ggml_hexagon_session * sess, + const struct ggml_tensor * src0, // values + const struct ggml_tensor * src1, // indices + const struct ggml_tensor * dst, // destination + struct htp_set_rows_kernel_params * kparams +) { + memset(kparams, 0, sizeof(*kparams)); + + const uint32_t nr = src0->ne[1]; + + kparams->n_threads = (std::min)((uint32_t)sess->n_threads, nr); + kparams->tasks_per_thread = (nr + kparams->n_threads - 1) / kparams->n_threads; + kparams->total_tasks = nr; + + kparams->div_ne11 = init_fastdiv_values(src1->ne[1]); + kparams->div_ne12 = init_fastdiv_values(src1->ne[2]); + kparams->div_tasks_per_thread = init_fastdiv_values(kparams->tasks_per_thread); + kparams->div_ne02 = init_fastdiv_values(src0->ne[2]); + + struct htp_set_rows_vtcm_layout vtcm_layout; + htp_set_rows_vtcm_layout_build(&vtcm_layout, dst->type, src0->ne[0], kparams->n_threads); + kparams->vtcm_size = vtcm_layout.total_bytes; } -static void ggml_hexagon_precompute_fused_ffn_params( +static void ggml_hexagon_precompute_fused_mmnx_params( const struct ggml_hexagon_session * sess, - const struct ggml_tensor * src0, // Wgate - const struct ggml_tensor * src1, // y + const struct ggml_tensor * src0, // W0 + const struct ggml_tensor * src1, // x + int32_t n_weights, struct htp_mm_kernel_params * kparams ) { memset(kparams, 0, sizeof(*kparams)); @@ -2737,7 +3913,7 @@ static void ggml_hexagon_precompute_fused_ffn_params( struct htp_mm_hvx_vtcm_layout L; htp_mm_hvx_vtcm_layout_build( &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0_row_size, src1_row_size, 0, d, false, false, true + 0, src0_row_size, src1_row_size, 0, d, false, true ); if (L.total_bytes <= sess->vtcm_size) { best_n_prefetch = d; @@ -2752,34 +3928,42 @@ static void ggml_hexagon_precompute_fused_ffn_params( // Test tiled first htp_mm_hvx_vtcm_layout_build( &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0_row_size, src1_row_size, 0, best_n_prefetch, false, false, true + 0, src0_row_size, src1_row_size, 0, best_n_prefetch, false, true ); if (try_tiled && L.total_bytes <= sess->vtcm_size) { kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW; kparams->vtcm_src0_size = L.src0_bytes; kparams->vtcm_src1_size = L.src1_bytes; - kparams->vtcm_src2_size = L.src2_bytes; kparams->vtcm_dst_size = L.dst_bytes; kparams->vtcm_size = L.total_bytes; kparams->n_prefetch = best_n_prefetch; + kparams->n_weights = n_weights; } else { kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT; size_t flat_src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); htp_mm_hvx_vtcm_layout_build( &L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0_row_size, flat_src1_row_size, 0, best_n_prefetch, false, false, true + 0, src0_row_size, flat_src1_row_size, 0, best_n_prefetch, false, true ); kparams->vtcm_src0_size = L.src0_bytes; kparams->vtcm_src1_size = L.src1_bytes; - kparams->vtcm_src2_size = L.src2_bytes; kparams->vtcm_dst_size = L.dst_bytes; kparams->vtcm_size = L.total_bytes; kparams->n_prefetch = best_n_prefetch; + kparams->n_weights = n_weights; } } +static bool ggml_hexagon_tensor_is_host(const struct ggml_hexagon_session * sess, const struct ggml_tensor * t) { + return t && t->buffer && t->buffer->buft == &sess->host_buffer_type; +} + +static bool ggml_hexagon_tensor_is_non_host(const struct ggml_hexagon_session * sess, const struct ggml_tensor * t) { + return t && t->buffer && t->buffer->buft != &sess->host_buffer_type; +} + static bool ggml_hexagon_supported_mul_mat(const struct ggml_hexagon_session * sess, const struct ggml_tensor * dst) { const struct ggml_tensor * src0 = dst->src[0]; const struct ggml_tensor * src1 = dst->src[1]; @@ -2811,9 +3995,8 @@ static bool ggml_hexagon_supported_mul_mat(const struct ggml_hexagon_session * s return false; // no broadcasting (for now) } - // src0 (weights) must be repacked - if (src0->buffer && !ggml_backend_buffer_is_hexagon_repack(src0->buffer)) { - return false; + if (!src0->buffer) { + sess->needs_repack.insert(src0); } break; @@ -2872,9 +4055,8 @@ static bool ggml_hexagon_supported_mul_mat_id(const struct ggml_hexagon_session return false; } - // src0 (weights) must be repacked - if (src0->buffer && !ggml_backend_buffer_is_hexagon_repack(src0->buffer)) { - return false; + if (!src0->buffer) { + sess->needs_repack.insert(src0); } break; @@ -2970,7 +4152,7 @@ static bool ggml_hexagon_supported_unary(const struct ggml_hexagon_session * ses if (dst->type != GGML_TYPE_F32) { return false; } - if (ggml_is_permuted(src0)) { + if (!ggml_is_contiguous_rows(src0)) { return false; } if (!ggml_are_same_shape(src0, dst)) { @@ -3114,7 +4296,11 @@ static bool ggml_hexagon_supported_softmax(const struct ggml_hexagon_session * s static bool ggml_hexagon_supported_set_rows(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) { const struct ggml_tensor * src0 = op->src[0]; // values const struct ggml_tensor * src1 = op->src[1]; // indices - const struct ggml_tensor * dst = op; + const struct ggml_tensor * dst = op->src[2] ? op->src[2] : op; + + if (dst->type == GGML_TYPE_Q8_0 && src0->ne[0] < 32) { + return false; + } if (src0->type != GGML_TYPE_F32) { return false; @@ -3124,7 +4310,7 @@ static bool ggml_hexagon_supported_set_rows(const struct ggml_hexagon_session * return false; } - if (dst->type != GGML_TYPE_F16) { + if (dst->type != GGML_TYPE_F32 && dst->type != GGML_TYPE_F16 && dst->type != GGML_TYPE_Q8_0) { return false; } @@ -3138,7 +4324,11 @@ static bool ggml_hexagon_supported_get_rows(const struct ggml_hexagon_session * const struct ggml_tensor * src1 = op->src[1]; // indices const struct ggml_tensor * dst = op; - if (src0->type != GGML_TYPE_F32) { + if (src0->type != GGML_TYPE_F32 && src0->ne[0] < 32) { + return false; + } + + if (src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_F16 && src0->type != GGML_TYPE_Q8_0) { return false; } @@ -3453,6 +4643,7 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) { case GGML_OP_CLAMP: return HTP_OP_CLAMP; case GGML_OP_SQR: return HTP_OP_SQR; case GGML_OP_SQRT: return HTP_OP_SQRT; + case GGML_OP_LOG: return HTP_OP_UNARY_LOG; case GGML_OP_SOFT_MAX: return HTP_OP_SOFTMAX; case GGML_OP_SSM_CONV: return HTP_OP_SSM_CONV; case GGML_OP_GATED_DELTA_NET: return HTP_OP_GATED_DELTA_NET; @@ -3476,6 +4667,7 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) { case GGML_UNARY_OP_EXP: return HTP_OP_UNARY_EXP; case GGML_UNARY_OP_SOFTPLUS: return HTP_OP_UNARY_SOFTPLUS; case GGML_UNARY_OP_TANH: return HTP_OP_UNARY_TANH; + case GGML_UNARY_OP_ABS: return HTP_OP_UNARY_ABS; default: break; } @@ -3521,133 +4713,22 @@ static bool is_mergeable_mul_mat(const ggml_tensor * t) { if (!t || t->op != GGML_OP_MUL_MAT) return false; if (t->src[1]->type != GGML_TYPE_F32) return false; return ggml_is_quantized(t->src[0]->type) && !mm_is_hmx_eligible(t); -} - -static bool is_mergeable_mul_mat_pair(const ggml_tensor * n1, const ggml_tensor * n2) { - if (!is_mergeable_mul_mat(n1) || !is_mergeable_mul_mat(n2)) { - return false; - } - if (n1->src[1] != n2->src[1]) { - return false; - } - if (n1->src[0]->ne[0] != n2->src[0]->ne[0] || - n1->src[0]->ne[1] != n2->src[0]->ne[1]) { - return false; - } - if (n1->src[0]->type != n2->src[0]->type) { - return false; - } - return true; -} - -static bool is_qkv_mergeable(const ggml_tensor * n_q, const ggml_tensor * n_k, const ggml_tensor * n_v) { - if (!is_mergeable_mul_mat(n_q) || !is_mergeable_mul_mat(n_k) || !is_mergeable_mul_mat(n_v)) { - return false; - } - if (n_q->src[1] != n_k->src[1] || n_q->src[1] != n_v->src[1]) { - return false; - } - if (n_q->src[0]->type != n_k->src[0]->type || n_q->src[0]->type != n_v->src[0]->type) { - return false; - } - if (n_k->src[0]->ne[0] != n_v->src[0]->ne[0] || - n_k->src[0]->ne[1] != n_v->src[0]->ne[1]) { - return false; - } - if (n_q->src[0]->ne[0] != n_k->src[0]->ne[0]) { - return false; - } - return true; -} - -static bool try_fuse_node(const ggml_hexagon_session * sess, const ggml_cgraph * graph, int & i, std::vector & nodes) { - if (!opt_opfusion) { - return false; - } - - ggml_tensor * n = graph->nodes[i]; - ggml_tensor * next_node = (i + 1 < graph->n_nodes) ? graph->nodes[i + 1] : nullptr; - - if (n->op == GGML_OP_RMS_NORM && next_node) { - if (next_node->op == GGML_OP_MUL && op_is_compute(next_node) && ggml_can_fuse(graph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL })) { - htp_opnode node(n, {}, HTP_OP_RMS_NORM_MUL); - node.add_fused(next_node); - - auto inputs = node.get_inputs(); - const struct ggml_tensor * src0 = inputs[0]; - const struct ggml_tensor * src1 = inputs.size() > 1 ? inputs[1] : nullptr; - ggml_hexagon_precompute_unary_params(sess, - node.opcode, src0, src1, node.dst(), - (struct htp_unary_kernel_params *)node.kernel_params - ); - - nodes.push_back(std::move(node)); - i++; // skip the fused MUL node - return true; - } - } - - if (is_mergeable_mul_mat(n)) { - ggml_tensor * n1 = (i + 1 < graph->n_nodes) ? graph->nodes[i + 1] : nullptr; - ggml_tensor * n2 = (i + 2 < graph->n_nodes) ? graph->nodes[i + 2] : nullptr; - if (is_qkv_mergeable(n, n1, n2)) { - struct htp_mm_kernel_params kparams; - ggml_hexagon_precompute_fused_qkv_params(sess, n1->src[0], n1->src[1], &kparams); - if ((size_t)kparams.vtcm_size <= sess->vtcm_size) { - // Reorder to KVQ: K (n1), V (n2), Q (n) - htp_opnode node(n1, {}, HTP_OP_MUL_MAT_QKV); - node.add_fused(n2, true); - node.add_fused(n, true); - memcpy(node.kernel_params, &kparams, sizeof(kparams)); - nodes.push_back(std::move(node)); - i += 2; - return true; - } else { - HEX_VERBOSE("ggml-hex: skip QKV fusion because VTCM needed (%d) > budget (%zu)\n", - kparams.vtcm_size, sess->vtcm_size); - } - } - if (is_mergeable_mul_mat_pair(n, n1)) { - struct htp_mm_kernel_params kparams; - ggml_hexagon_precompute_fused_ffn_params(sess, n->src[0], n->src[1], &kparams); - if ((size_t)kparams.vtcm_size <= sess->vtcm_size) { - htp_opnode node(n, {}, HTP_OP_MUL_MAT_FFN); - node.add_fused(n1, true); - memcpy(node.kernel_params, &kparams, sizeof(kparams)); - nodes.push_back(std::move(node)); - i += 1; - return true; - } else { - HEX_VERBOSE("ggml-hex: skip FFN fusion because VTCM needed (%d) > budget (%zu)\n", - kparams.vtcm_size, sess->vtcm_size); - } - } - } +} - if (n->op == GGML_OP_MUL_MAT && next_node) { - if (next_node->op == GGML_OP_ADD && op_is_compute(next_node) && ggml_can_fuse(graph, i, { GGML_OP_MUL_MAT, GGML_OP_ADD })) { - if (next_node->src[0] == n || next_node->src[1] == n) { - const struct ggml_tensor * src2 = (next_node->src[0] == n) ? next_node->src[1] : next_node->src[0]; - struct htp_mm_kernel_params kparams; - ggml_hexagon_precompute_fused_matmul_add_params(sess, n->src[0], n->src[1], src2, next_node, &kparams); - const int src1_nrows = n->src[1]->ne[1] * n->src[1]->ne[2] * n->src[1]->ne[3]; - const bool can_fuse = (kparams.n_hmx > 0) || (src1_nrows == 1); - if (can_fuse && (size_t)kparams.vtcm_size <= sess->vtcm_size) { - htp_opnode node(n, {}, HTP_OP_MUL_MAT_ADD); - node.add_fused(next_node); - memcpy(node.kernel_params, &kparams, sizeof(kparams)); - nodes.push_back(std::move(node)); - i += 1; - return true; - } else if (can_fuse) { - HEX_VERBOSE("ggml-hex: skip MUL_MAT_ADD fusion because VTCM needed (%d) > budget (%zu)\n", - kparams.vtcm_size, sess->vtcm_size); - } - } - } +static bool is_mergeable_mul_mat_pair(const ggml_tensor * n1, const ggml_tensor * n2) { + if (!is_mergeable_mul_mat(n1) || !is_mergeable_mul_mat(n2)) { + return false; } - - return false; + if (n1->src[1] != n2->src[1]) { + return false; + } + if (n1->src[0]->ne[0] != n2->src[0]->ne[0]) { + return false; + } + if (n1->src[0]->type != n2->src[0]->type) { + return false; + } + return true; } static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, ggml_cgraph * graph) { @@ -3659,24 +4740,34 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg std::vector computed_nodes; // Check for cache hit - bool cache_hit = (graph->uid != 0 && sess->cached_graph.uid == graph->uid); + bool cache_hit = (graph->uid != 0 && sess->cached_uid == graph->uid); if (cache_hit) { - nodes_ptr = &sess->cached_graph.htp_nodes; + nodes_ptr = &sess->cached_nodes; } else { + // Tag fusable tensors in graph + for (int i = 0; i < graph->n_nodes; i++) { + auto * extra = (ggml_hexagon_tensor_extra *) graph->nodes[i]->extra; + if (!extra) continue; + + if (graph->nodes[i]->op == GGML_OP_RMS_NORM && ggml_can_fuse(graph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL })) { + extra->flags |= GGML_HEXAGON_TENSOR_FUSEABLE; + } else if (graph->nodes[i]->op == GGML_OP_MUL_MAT) { + if ((i + 1 < graph->n_nodes && graph->nodes[i + 1]->op == GGML_OP_ADD && ggml_can_fuse(graph, i, { GGML_OP_MUL_MAT, GGML_OP_ADD })) || + ggml_node_has_n_uses(graph, i, 1)) { + extra->flags |= GGML_HEXAGON_TENSOR_FUSEABLE; + } + } + } + computed_nodes.reserve(graph->n_nodes); - // Fuse and finalize for (int i = 0; i < graph->n_nodes; ++i) { ggml_tensor * n = graph->nodes[i]; if (!op_is_compute(n)) { continue; } - if (try_fuse_node(sess, graph, i, computed_nodes)) { - continue; - } - - htp_opnode node(n, {}, HTP_OP_INVALID); + htp_opnode node(HTP_OP_INVALID, n); node.opcode = op_remap_to_htp(n); if (node.opcode == HTP_OP_MUL_MAT || node.opcode == HTP_OP_MUL_MAT_ID) { ggml_hexagon_precompute_matmul_params(sess, @@ -3696,29 +4787,34 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg node.opcode, src0, src1, node.dst(), (struct htp_unary_kernel_params *)node.kernel_params ); + } else if (node.opcode == HTP_OP_GET_ROWS) { + ggml_hexagon_precompute_get_rows_params(sess, + node.node->src[0], node.node->src[1], node.dst(), + (struct htp_get_rows_kernel_params *)node.kernel_params + ); + } else if (node.opcode == HTP_OP_SET_ROWS) { + ggml_hexagon_precompute_set_rows_params(sess, + node.node->src[0], node.node->src[1], node.dst(), + (struct htp_set_rows_kernel_params *)node.kernel_params + ); } computed_nodes.push_back(std::move(node)); } if (graph->uid != 0) { - sess->cached_graph.uid = graph->uid; - sess->cached_graph.htp_nodes = std::move(computed_nodes); - nodes_ptr = &sess->cached_graph.htp_nodes; + sess->cached_uid = graph->uid; + sess->cached_nodes = std::move(computed_nodes); + nodes_ptr = &sess->cached_nodes; } else { nodes_ptr = &computed_nodes; } } // Queue and execute - if (opt_opstage & HTP_OPSTAGE_QUEUE) { - for (const auto & node : *nodes_ptr) { - sess->enqueue_op(node); - } + for (const auto & node : *nodes_ptr) { + sess->enqueue_op(node); } - // Wait until all pending ops complete - sess->flush(); - return GGML_STATUS_SUCCESS; } @@ -3731,6 +4827,106 @@ static void ggml_backend_hexagon_synchronize(ggml_backend_t backend) { sess->flush(); } +enum ggml_hexagon_mem_range_type { + HEXAGON_MEM_RANGE_TYPE_SRC, + HEXAGON_MEM_RANGE_TYPE_DST, +}; + +struct ggml_hexagon_mem_range { + uint64_t pb; + uint64_t p0; + uint64_t p1; + ggml_hexagon_mem_range_type pt; +}; + +struct ggml_hexagon_mem_ranges { + std::vector ranges; + + void reset() { + ranges.clear(); + } + + void add(const ggml_hexagon_mem_range & mr) { + ranges.push_back(mr); + } + + bool check(const ggml_hexagon_mem_range & mr) const { + for (const auto & cmp : ranges) { + if (mr.pb != cmp.pb) { + continue; + } + if (mr.pt == HEXAGON_MEM_RANGE_TYPE_SRC && cmp.pt == HEXAGON_MEM_RANGE_TYPE_SRC) { + continue; + } + if (mr.p0 < cmp.p1 && mr.p1 > cmp.p0) { + return false; + } + } + return true; + } +}; + +static ggml_hexagon_mem_range ggml_hexagon_mem_range_from_tensor(const ggml_tensor * tensor, ggml_hexagon_mem_range_type pt) { + const ggml_tensor * base = tensor->view_src ? tensor->view_src : tensor; + ggml_hexagon_mem_range mr; + if (tensor->buffer) { + mr = { + /*.pb =*/ (uint64_t) tensor->buffer, + /*.p0 =*/ (uint64_t) tensor->data, + /*.p1 =*/ (uint64_t) tensor->data + ggml_backend_buft_get_alloc_size(tensor->buffer->buft, tensor), + /*.pt =*/ pt, + }; + } else { + mr = { + /*.pb =*/ (uint64_t) base, + /*.p0 =*/ 0, + /*.p1 =*/ 1024, + /*.pt =*/ pt, + }; + } + return mr; +} + +static void ggml_hexagon_mem_ranges_add_node(ggml_hexagon_mem_ranges & mrs, const htp_opnode & node) { + if (node.is_empty()) return; + + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (node.node->src[i]) { + mrs.add(ggml_hexagon_mem_range_from_tensor(node.node->src[i], HEXAGON_MEM_RANGE_TYPE_SRC)); + } + } + for (const auto * fused : node.fused) { + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (fused->src[i]) { + mrs.add(ggml_hexagon_mem_range_from_tensor(fused->src[i], HEXAGON_MEM_RANGE_TYPE_SRC)); + } + } + } + mrs.add(ggml_hexagon_mem_range_from_tensor(node.dst(), HEXAGON_MEM_RANGE_TYPE_DST)); +} + +static bool ggml_hexagon_mem_ranges_check_node(const ggml_hexagon_mem_ranges & mrs, const htp_opnode & node) { + if (node.is_empty()) return true; + + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (node.node->src[i]) { + if (!mrs.check(ggml_hexagon_mem_range_from_tensor(node.node->src[i], HEXAGON_MEM_RANGE_TYPE_SRC))) { + return false; + } + } + } + for (const auto * fused : node.fused) { + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (fused->src[i]) { + if (!mrs.check(ggml_hexagon_mem_range_from_tensor(fused->src[i], HEXAGON_MEM_RANGE_TYPE_SRC))) { + return false; + } + } + } + } + return mrs.check(ggml_hexagon_mem_range_from_tensor(node.dst(), HEXAGON_MEM_RANGE_TYPE_DST)); +} + static std::vector ggml_hexagon_graph_optimize_reorder(const std::vector & nodes) { const int n = nodes.size(); @@ -3739,28 +4935,32 @@ static std::vector ggml_hexagon_graph_optimize_reorder(const std::vector used(n, false); - // The main goal here is to stack the MUL_MAT ops with the same src1 input. - // This allows use to reuse dynamically quantized src1 in VTCM. + ggml_hexagon_mem_ranges mrs; - // TODO: the current version might do incorrect reordering in cases where quantized src0 - // input is an output of another Op. + // The main goal here is to stack the MUL_MAT ops with the same src1 input. + // This allows us to reuse dynamically quantized src1 in VTCM. for (int i0 = 0; i0 < n; i0++) { if (used[i0]) { continue; } - res.push_back(i0); - const auto & node0 = nodes[i0]; if (!node0.stackable()) { + res.push_back(i0); + used[i0] = true; continue; } // that many nodes forward to search for stackable nodes that can reuse VTCM constexpr int N_FORWARD = 16; + std::vector stack; + stack.push_back(i0); + + mrs.reset(); + for (int i1 = i0 + 1; i1 < i0 + N_FORWARD && i1 < n; i1++) { if (used[i1]) { continue; @@ -3768,11 +4968,17 @@ static std::vector ggml_hexagon_graph_optimize_reorder(const std::vector nodes; nodes.reserve(gf->n_nodes); - // fuse nodes: - // we don't want to make reorders that break fusing, so we first pack all fusable tensors - // and perform the reorder over the fused nodes. after the reorder is done, we unfuse + // Pack nodes for reordering for (int i = 0; i < n; i++) { - htp_opnode node = { - /*.node =*/gf->nodes[i], - /*.fused =*/{}, - }; + htp_opnode node(HTP_OP_INVALID, gf->nodes[i]); // fuse only ops that start with these operations // can be expanded when needed @@ -3856,22 +5057,200 @@ static void ggml_backend_hexagon_graph_optimize(ggml_backend_t backend, ggml_cgr GGML_UNUSED(backend); } +static bool ggml_hexagon_cpy_tensor_async_phys(ggml_backend_t backend_src, ggml_backend_t backend_dst, const ggml_tensor * src, ggml_tensor * dst) { + auto sess_src = static_cast(backend_src->context); + auto sess_dst = static_cast(backend_dst->context); + auto sbuf_dst = (ggml_hexagon_shared_buffer *) dst->buffer->context; + + if (sess_dst->fence_seq == 0) sess_dst->fence_seq = 1; + uint32_t fence_seq = sess_dst->fence_seq++; + if (sess_dst->fence_seq == 0) sess_dst->fence_seq = 1; + + volatile uint32_t * fence = (volatile uint32_t *) sbuf_dst->alloc_fence(); + + HEX_VERBOSE("ggml-hex: %s cpy-tensor-async %s -> %s size %zu : seq %u\n", + sess_dst->name.c_str(), src->name, dst->name, ggml_nbytes(src), fence_seq); + + // dummy extra (must be static) + static ggml_hexagon_tensor_extra fence_extra { {}, 0, GGML_HEXAGON_TENSOR_FENCE }; + + ggml_tensor fence_tensor {}; + fence_tensor.buffer = dst->buffer; + fence_tensor.extra = &fence_extra; + fence_tensor.data = (void *) fence; + fence_tensor.type = GGML_TYPE_I32; + fence_tensor.ne[0] = 1; + fence_tensor.ne[1] = 1; + fence_tensor.ne[2] = 1; + fence_tensor.ne[3] = 1; + fence_tensor.nb[0] = sizeof(int32_t); + fence_tensor.nb[1] = sizeof(int32_t); + fence_tensor.nb[2] = sizeof(int32_t); + fence_tensor.nb[3] = sizeof(int32_t); + fence_tensor.op = GGML_OP_NONE; + + sess_src->enqueue_cpy(src, dst, &fence_tensor, fence_seq); + sess_dst->enqueue_fence(&fence_tensor, fence_seq); + + sess_dst->add_sync_peer(sess_src); + + return true; +} + +static bool ggml_hexagon_cpy_tensor_async_virt(ggml_backend_t backend_src, ggml_backend_t backend_dst, const ggml_tensor * src, ggml_tensor * dst) { + auto sess_src = static_cast(backend_src->context); + auto sess_dst = static_cast(backend_dst->context); + auto sbuf_dst = (ggml_hexagon_shared_buffer *) dst->buffer->context; + + if (!sess_src->clone_buffer(sbuf_dst)) { return false; } + + HEX_VERBOSE("ggml-hex: %s cpy-tensor-async %s -> %s size %zu\n", + sess_dst->name.c_str(), src->name, dst->name, ggml_nbytes(src)); + + sess_src->enqueue_cpy(src, dst); + sess_src->flush(true); + + return true; +} + +static bool ggml_backend_hexagon_cpy_tensor_async(ggml_backend_t backend_src, ggml_backend_t backend_dst, const ggml_tensor * src, ggml_tensor * dst) { + if (!ggml_backend_is_hexagon(backend_src) || !ggml_backend_is_hexagon(backend_dst)) { + return false; + } + + *(ggml_hexagon_tensor_extra *) dst->extra = *(const ggml_hexagon_tensor_extra *) src->extra; + + auto sess_src = static_cast(backend_src->context); + auto sess_dst = static_cast(backend_dst->context); + + if (sess_src == sess_dst) { + HEX_VERBOSE("ggml-hex: %s cpy-tensor-async %s -> %s size %zu\n", sess_dst->name.c_str(), src->name, dst->name, ggml_nbytes(src)); + sess_src->enqueue_cpy(src, dst); + sess_src->flush_batch(); + return true; + } + + if (sess_src->phys_idx != sess_dst->phys_idx) + return ggml_hexagon_cpy_tensor_async_phys(backend_src, backend_dst, src, dst); + + return ggml_hexagon_cpy_tensor_async_virt(backend_src, backend_dst, src, dst); +} + +static ggml_backend_event_t ggml_backend_hexagon_device_event_new(ggml_backend_dev_t dev) { + ggml_hexagon_event * hex_event = new ggml_hexagon_event(); + HEX_VERBOSE("ggml-hex: %s event-new : event %p\n", ggml_backend_dev_name(dev), (void *)hex_event); + + return new ggml_backend_event { + /* .device = */ dev, + /* .context = */ hex_event, + }; +} + +static void ggml_backend_hexagon_device_event_free(ggml_backend_dev_t dev, ggml_backend_event_t event) { + GGML_UNUSED(dev); + + if (event == nullptr) { + return; + } + + ggml_hexagon_event * hex_event = (ggml_hexagon_event *)event->context; + HEX_VERBOSE("ggml-hex: %s event-free : event %p\n", ggml_backend_dev_name(dev), (void *)hex_event); + delete hex_event; + delete event; +} + +static void ggml_backend_hexagon_device_event_synchronize(ggml_backend_dev_t dev, ggml_backend_event_t event) { + GGML_UNUSED(dev); + + ggml_hexagon_event * hex_event = (ggml_hexagon_event *)event->context; + HEX_VERBOSE("ggml-hex: %s event-synchronize : event %p seq %llu\n", + ggml_backend_dev_name(dev), (void *)hex_event, (unsigned long long)hex_event->seq); + if (hex_event->sess != nullptr) { + hex_event->sess->wait_event(hex_event->seq); + } +} + +static void ggml_backend_hexagon_event_record(ggml_backend_t backend, ggml_backend_event_t event) { + auto sess = static_cast(backend->context); + ggml_hexagon_event * hex_event = (ggml_hexagon_event *)event->context; + + hex_event->sess = sess; + hex_event->seq = sess->record_event(); + HEX_VERBOSE("ggml-hex: %s event-record : event %p seq %llu\n", + sess->c_name(), (void *)hex_event, (unsigned long long)hex_event->seq); +} + +static void ggml_backend_hexagon_event_wait(ggml_backend_t backend, ggml_backend_event_t event) { + GGML_UNUSED(backend); + + ggml_hexagon_event * hex_event = (ggml_hexagon_event *)event->context; + if (hex_event->sess != nullptr) { + HEX_VERBOSE("ggml-hex: %s event-wait : event %p seq %llu\n", + hex_event->sess->c_name(), (void *)hex_event, (unsigned long long)hex_event->seq); + hex_event->sess->wait_event(hex_event->seq); + } +} + +static void ggml_backend_hexagon_set_tensor_async(ggml_backend_t backend, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size) { + auto sess = static_cast(backend->context); + HEX_VERBOSE("ggml-hex: %s set-tensor-async %s : data %p offset %zu size %zu usage %d\n", + sess->c_name(), tensor->name, data, offset, size, tensor->buffer ? (int) tensor->buffer->usage : -1); + ggml_backend_tensor_set(tensor, data, offset, size); +} + +static void ggml_backend_hexagon_get_tensor_async(ggml_backend_t backend, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size) { + auto sess = static_cast(backend->context); + HEX_VERBOSE("ggml-hex: %s get-tensor-async %s : data %p offset %zu size %zu usage %d\n", + sess->c_name(), tensor->name, data, offset, size, tensor->buffer ? (int) tensor->buffer->usage : -1); + sess->flush(true); + ggml_backend_tensor_get(tensor, data, offset, size); +} + +static void ggml_backend_hexagon_set_tensor_2d_async(ggml_backend_t backend, + 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) { + auto sess = static_cast(backend->context); + HEX_VERBOSE("ggml-hex: %s set-tensor-2d-async %s : data %p offset %zu size %zu n_copies %zu stride_tensor %zu stride_data %zu usage %d\n", + sess->c_name(), tensor->name, data, offset, size, n_copies, stride_tensor, stride_data, tensor->buffer ? (int) tensor->buffer->usage : -1); + ggml_backend_tensor_set_2d(tensor, data, offset, size, n_copies, stride_tensor, stride_data); +} + +static void ggml_backend_hexagon_get_tensor_2d_async(ggml_backend_t backend, + 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) { + auto sess = static_cast(backend->context); + HEX_VERBOSE("ggml-hex: %s get-tensor-2d-async %s : data %p offset %zu size %zu n_copies %zu stride_tensor %zu stride_data %zu usage %d\n", + sess->c_name(), tensor->name, data, offset, size, n_copies, stride_tensor, stride_data, tensor->buffer ? (int) tensor->buffer->usage : -1); + sess->flush(true); + ggml_backend_tensor_get_2d(tensor, data, offset, size, n_copies, stride_tensor, stride_data); +} + static struct ggml_backend_i hexagon_backend_i = { /* .get_name = */ ggml_backend_hexagon_name, /* .free = */ ggml_backend_hexagon_free, - /* .set_tensor_async = */ NULL, - /* .get_tensor_async = */ NULL, - /* .set_tensor_2d_async = */ NULL, - /* .get_tensor_2d_async = */ NULL, - /* .cpy_tensor_async = */ NULL, + /* .set_tensor_async = */ ggml_backend_hexagon_set_tensor_async, + /* .get_tensor_async = */ ggml_backend_hexagon_get_tensor_async, + /* .set_tensor_2d_async = */ ggml_backend_hexagon_set_tensor_2d_async, + /* .get_tensor_2d_async = */ ggml_backend_hexagon_get_tensor_2d_async, + /* .cpy_tensor_async = */ ggml_backend_hexagon_cpy_tensor_async, /* .synchronize = */ ggml_backend_hexagon_synchronize, /* .graph_plan_create = */ NULL, /* .graph_plan_free = */ NULL, /* .graph_plan_update = */ NULL, /* .graph_plan_compute = */ NULL, /* .graph_compute = */ ggml_backend_hexagon_graph_compute, - /* .event_record = */ NULL, - /* .event_wait = */ NULL, + /* .event_record = */ ggml_backend_hexagon_event_record, + /* .event_wait = */ ggml_backend_hexagon_event_wait, /* .graph_optimize = */ ggml_backend_hexagon_graph_optimize, }; @@ -3932,9 +5311,9 @@ static void ggml_backend_hexagon_device_get_props(ggml_backend_dev_t dev, struct ggml_backend_hexagon_device_get_memory(dev, &props->memory_free, &props->memory_total); props->caps = { /* .async = */ true, - /* .host_buffer = */ (bool) opt_hostbuf, + /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ false, - /* .events = */ false, + /* .events = */ true, /* .mmap_support = */ false, }; } @@ -3944,32 +5323,12 @@ static ggml_backend_buffer_type_t ggml_backend_hexagon_device_get_buffer_type(gg return &sess->buffer_type; } -static ggml_backend_buffer_type_t ggml_backend_hexagon_device_get_repack_buffer_type(ggml_backend_dev_t dev) { - auto sess = static_cast(dev->context); - return &sess->repack_buffer_type; -} - -static bool ggml_hexagon_supported_buffer(ggml_hexagon_session *sess, const struct ggml_tensor * t) { - if (t && t->buffer) { - if (ggml_backend_buffer_is_hexagon(t->buffer) == false) return false; // not our buffer - if (ggml_backend_hexagon_buffer_get_sess(t->buffer) != sess) return false; // wrong session - } - return true; -} - -static bool ggml_hexagon_supported_buffers(ggml_hexagon_session *sess, const struct ggml_tensor * t) { - // all srcs & dsts must be mapped to the same session - if (!ggml_hexagon_supported_buffer(sess, t)) { - return false; - } - - for (int i = 0; i < GGML_MAX_SRC; i++) { - if (!ggml_hexagon_supported_buffer(sess, t->src[i])) { - return false; - } +static ggml_backend_buffer_type_t ggml_backend_hexagon_device_get_host_buffer_type(ggml_backend_dev_t dev) { + if (!opt_hostbuf) { + return NULL; } - - return true; + auto sess = static_cast(dev->context); + return &sess->host_buffer_type; } static bool ggml_hexagon_supported_cpy(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) { @@ -4067,12 +5426,6 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons return false; } - // all srcs & dsts must be mapped to the same session - if (!ggml_hexagon_supported_buffers(sess, op)) { - ggml_hexagon_dump_op_supp(sess->name, op, false); - return false; - } - bool supp = false; switch (op->op) { case GGML_OP_NONE: @@ -4112,6 +5465,7 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons case GGML_OP_SQR: case GGML_OP_SQRT: + case GGML_OP_LOG: supp = ggml_hexagon_supported_unary(sess, op); break; @@ -4130,6 +5484,7 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons case GGML_UNARY_OP_SIGMOID: case GGML_UNARY_OP_SOFTPLUS: case GGML_UNARY_OP_TANH: + case GGML_UNARY_OP_ABS: case GGML_UNARY_OP_SILU: case GGML_UNARY_OP_GELU: case GGML_UNARY_OP_GELU_QUICK: @@ -4233,31 +5588,20 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons } static bool ggml_backend_hexagon_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { - if (buft->iface.get_alignment != ggml_backend_hexagon_buffer_type_get_alignment) { - return false; - } - - auto s0 = static_cast(dev->context); - auto s1 = static_cast(buft->context)->sess; - - // Need session/domain-id for buffers to be compatible - bool supp = (s0->session_id == s1->session_id); + auto sess = static_cast(dev->context); - HEX_VERBOSE("ggml-hex: %s device-supports-buft %s (%d)\n", s0->name.c_str(), s1->name.c_str(), (int) supp); + // Technically we can clone hexagon buffers from any session but for some reason the output is garbled with layer-split, + // tensor-split works correctly, so it needs mode debugging and investigation. For now accept only our own buffers. +#if 0 + bool supp = (buft->iface.get_alignment == ggml_backend_hexagon_buffer_type_get_alignment); +#else + bool supp = (buft == &sess->host_buffer_type) || (buft == &sess->buffer_type); +#endif + HEX_VERBOSE("ggml-hex: %s device-supports-buft %s %s\n", sess->name.c_str(), ggml_backend_buft_name(buft), supp ? "yes" : "no"); return supp; } -static ggml_backend_buffer_type_t * ggml_backend_hexagon_device_get_extra_buffers_type(ggml_backend_dev_t dev) { - auto s0 = static_cast(dev->context); - HEX_VERBOSE("ggml-hex: device-get-extra-buft : %s \n", s0->name.c_str()); - - static ggml_backend_buffer_type_t bufts[2]; - bufts[0] = ggml_backend_hexagon_device_get_repack_buffer_type(dev); - bufts[1] = NULL; - return bufts; -} - static const struct ggml_backend_device_i ggml_backend_hexagon_device_i = { /* .get_name = */ ggml_backend_hexagon_device_get_name, /* .get_description = */ ggml_backend_hexagon_device_get_description, @@ -4266,27 +5610,18 @@ static const struct ggml_backend_device_i ggml_backend_hexagon_device_i = { /* .get_props = */ ggml_backend_hexagon_device_get_props, /* .init_backend = */ ggml_backend_hexagon_device_init, /* .get_buffer_type = */ ggml_backend_hexagon_device_get_buffer_type, - /* .get_host_buffer_type = */ NULL, // ggml_backend_hexagon_device_get_host_buffer_type, + /* .get_host_buffer_type = */ ggml_backend_hexagon_device_get_host_buffer_type, /* .buffer_from_host_ptr = */ NULL, // ggml_backend_hexagon_device_buffer_from_ptr, /* .supports_op = */ ggml_backend_hexagon_device_supports_op, /* .supports_buft = */ ggml_backend_hexagon_device_supports_buft, /* .offload_op = */ NULL, // ggml_backend_hexagon_device_offload_op, - /* .event_new = */ NULL, - /* .event_free = */ NULL, - /* .event_synchronize = */ NULL, + /* .event_new = */ ggml_backend_hexagon_device_event_new, + /* .event_free = */ ggml_backend_hexagon_device_event_free, + /* .event_synchronize = */ ggml_backend_hexagon_device_event_synchronize, }; //** backend registry -#define GGML_HEXAGON_MAX_SESSIONS 16 - -struct ggml_hexagon_registry { - ggml_hexagon_registry(ggml_backend_reg_t reg); - ~ggml_hexagon_registry(); - - ggml_backend_device devices[GGML_HEXAGON_MAX_SESSIONS]; -}; - ggml_hexagon_registry::ggml_hexagon_registry(ggml_backend_reg_t reg) { GGML_LOG_INFO("ggml-hex: Hexagon backend (experimental) : allocating new registry : ndev %zu\n", opt_ndev); @@ -4303,6 +5638,7 @@ ggml_hexagon_registry::ggml_hexagon_registry(ggml_backend_reg_t reg) { devices[i].context = nullptr; } } + } ggml_hexagon_registry::~ggml_hexagon_registry() { @@ -4335,14 +5671,129 @@ static ggml_backend_dev_t ggml_backend_hexagon_reg_get_device(ggml_backend_reg_t return &hreg->devices[index]; } -static void * ggml_backend_hexagon_get_proc_address(ggml_backend_reg_t reg, const char * name) { - if (strcmp(name, "ggml_backend_dev_get_extra_bufts") == 0 && opt_hostbuf) { - ggml_backend_dev_get_extra_bufts_t fct = ggml_backend_hexagon_device_get_extra_buffers_type; - return (void *) fct; +// ** communication context for tensor-split allreduce + +static void * ggml_backend_hexagon_comm_init(ggml_backend_t * backends, size_t n_backends) { + if (n_backends < 2 || n_backends > 4) { + return nullptr; } - return NULL; + for (size_t i = 0; i < n_backends; ++i) { + if (!ggml_backend_is_hexagon(backends[i])) { + return nullptr; + } + } + + auto * ctx = new ggml_backend_hexagon_comm_context(); + ctx->backends.assign(backends, backends + n_backends); + ctx->n_backends = n_backends; + ctx->fence_seq = (((uintptr_t) ctx) & 0xFFFF) | 1; + + return ctx; +} + +static void ggml_backend_hexagon_comm_free(void * comm_ctx_v) { + if (!comm_ctx_v) return; + delete static_cast(comm_ctx_v); +} + +static bool ggml_backend_hexagon_comm_allreduce_tensor(void * comm_ctx_v, struct ggml_tensor ** tensors) { + if (opt_ar_select == 0 || !comm_ctx_v) return false; + auto * comm_ctx = static_cast(comm_ctx_v); + const size_t n_backends = comm_ctx->n_backends; + + if (n_backends < 2 || n_backends > 4) return false; + + for (size_t i = 0; i < n_backends; i++) { + if (!tensors[i] || !tensors[i]->buffer || !ggml_backend_buffer_is_hexagon(tensors[i]->buffer)) { + return false; + } + if (tensors[i]->type != tensors[0]->type) { + return false; + } + if (!ggml_is_contiguous(tensors[i])) { + return false; + } + if (ggml_nelements(tensors[i]) != ggml_nelements(tensors[0])) { + return false; + } + } + + if (tensors[0]->type != GGML_TYPE_F16 && tensors[0]->type != GGML_TYPE_F32) { + return false; + } + + for (size_t r = 0; r < n_backends; r++) { + auto sess = static_cast(comm_ctx->backends[r]->context); + struct htp_allreduce_kernel_params kparams; + if (!ggml_hexagon_precompute_allreduce_params(sess, tensors[r], (uint32_t) r, (uint32_t) n_backends, false, false, &kparams)) { + return false; + } + } + + if (comm_ctx->fence_seq == 0) comm_ctx->fence_seq = 1; + uint32_t fence_seq_entry = comm_ctx->fence_seq++; + if (comm_ctx->fence_seq == 0) comm_ctx->fence_seq = 1; + uint32_t fence_seq_exit = comm_ctx->fence_seq++; + if (comm_ctx->fence_seq == 0) comm_ctx->fence_seq = 1; + + volatile uint32_t * fences[GGML_HEXAGON_MAX_SESSIONS]; + for (size_t i = 0; i < n_backends; i++) { + auto sbuf = (ggml_hexagon_shared_buffer *) tensors[i]->buffer->context; + fences[i] = (volatile uint32_t *) sbuf->alloc_fence(); + } + + static ggml_hexagon_tensor_extra fence_extra { {}, 0, GGML_HEXAGON_TENSOR_FENCE }; + ggml_tensor fence_tensors[GGML_HEXAGON_MAX_SESSIONS]; + for (size_t i = 0; i < n_backends; i++) { + fence_tensors[i] = {}; + fence_tensors[i].buffer = tensors[i]->buffer; + fence_tensors[i].extra = &fence_extra; + fence_tensors[i].data = (void *) fences[i]; + fence_tensors[i].type = GGML_TYPE_I32; + fence_tensors[i].ne[0] = 4; + fence_tensors[i].ne[1] = 1; + fence_tensors[i].ne[2] = 1; + fence_tensors[i].ne[3] = 1; + fence_tensors[i].nb[0] = sizeof(int32_t); + fence_tensors[i].nb[1] = sizeof(int32_t); + fence_tensors[i].nb[2] = sizeof(int32_t); + fence_tensors[i].nb[3] = sizeof(int32_t); + fence_tensors[i].op = GGML_OP_NONE; + } + + std::vector data_tensors(n_backends); + std::vector sync_tensors(n_backends); + for (size_t i = 0; i < n_backends; i++) { + data_tensors[i] = tensors[i]; + sync_tensors[i] = &fence_tensors[i]; + } + + for (size_t r = 0; r < n_backends; r++) { + auto sess = static_cast(comm_ctx->backends[r]->context); + sess->enqueue_allreduce(tensors[r], data_tensors, sync_tensors, (uint32_t) r, (uint32_t) n_backends, fence_seq_entry, fence_seq_exit); + for (size_t j = 0; j < n_backends; j++) { + if (r != j) { + sess->add_sync_peer(static_cast(comm_ctx->backends[j]->context)); + } + } + } + + return true; +} + +static void * ggml_backend_hexagon_get_proc_address(ggml_backend_reg_t reg, const char * name) { GGML_UNUSED(reg); + if (strcmp(name, "ggml_backend_comm_init") == 0) { + return (void *) ggml_backend_hexagon_comm_init; + } + if (strcmp(name, "ggml_backend_comm_free") == 0) { + return (void *) ggml_backend_hexagon_comm_free; + } + if (strcmp(name, "ggml_backend_comm_allreduce_tensor") == 0) { + return (void *) ggml_backend_hexagon_comm_allreduce_tensor; + } + return NULL; } template std::vector str_to_vec(const char* str) { @@ -4379,8 +5830,6 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) { "please update hexagon_type to match ggml_type"); const char * str_verbose = getenv("GGML_HEXAGON_VERBOSE"); - const char * str_hostbuf = getenv("GGML_HEXAGON_HOSTBUF"); - const char * str_opstage = getenv("GGML_HEXAGON_OPSTAGE"); const char * str_opbatch = getenv("GGML_HEXAGON_OPBATCH"); const char * str_opqueue = getenv("GGML_HEXAGON_OPQUEUE"); const char * str_oppoll = getenv("GGML_HEXAGON_OPPOLL"); @@ -4389,15 +5838,16 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) { const char * str_profile = getenv("GGML_HEXAGON_PROFILE"); const char * str_etm = getenv("GGML_HEXAGON_ETM"); const char * str_nhvx = getenv("GGML_HEXAGON_NHVX"); - const char * str_use_hmx = getenv("GGML_HEXAGON_USE_HMX"); const char * str_nhmx = getenv("GGML_HEXAGON_NHMX"); const char * str_mm_select = getenv("GGML_HEXAGON_MM_SELECT"); const char * str_fa_select = getenv("GGML_HEXAGON_FA_SELECT"); + const char * str_ar_select = getenv("GGML_HEXAGON_AR_SELECT"); const char * str_ndev = getenv("GGML_HEXAGON_NDEV"); const char * str_arch = getenv("GGML_HEXAGON_ARCH"); const char * str_vmem = getenv("GGML_HEXAGON_VMEM"); const char * str_mbuf = getenv("GGML_HEXAGON_MBUF"); const char * str_optrace = getenv("GGML_HEXAGON_OPTRACE"); + const char * str_hostbuf = getenv("GGML_HEXAGON_HOSTBUF"); // Init Arch first since it affects other defaults if (!str_arch) { @@ -4430,8 +5880,6 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) { opt_opfilter = str_opfilter ? new std::regex(str_opfilter, RE_ICASE) : NULL; opt_verbose = str_verbose ? atoi(str_verbose) : 0; - opt_hostbuf = str_hostbuf ? atoi(str_hostbuf) : opt_hostbuf; - opt_opstage = str_opstage ? strtoul(str_opstage, NULL, 0) : opt_opstage; opt_opbatch = str_opbatch ? strtoul(str_opbatch, NULL, 0) : opt_opbatch; opt_opqueue = str_opqueue ? strtoul(str_opqueue, NULL, 0) : opt_opqueue; opt_optrace = str_optrace ? strtoul(str_optrace, NULL, 0) : (opt_opbatch * 256); @@ -4440,16 +5888,90 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) { opt_profile = str_profile ? atoi(str_profile) : 0; opt_etm = str_etm ? atoi(str_etm) : 0; opt_nhvx = str_nhvx ? strtoul(str_nhvx, NULL, 0) : opt_nhvx; - opt_nhmx = str_nhmx ? atoi(str_nhmx) : (str_use_hmx ? atoi(str_use_hmx) : opt_nhmx); + opt_nhmx = str_nhmx ? atoi(str_nhmx) : opt_nhmx; opt_mm_select = str_mm_select ? atoi(str_mm_select) : opt_mm_select; opt_fa_select = str_fa_select ? atoi(str_fa_select) : opt_fa_select; - opt_ndev = str_ndev ? strtoul(str_ndev, NULL, 0) : opt_ndev; - opt_hostbuf = str_hostbuf ? atoi(str_hostbuf) : opt_hostbuf; + opt_ar_select = str_ar_select ? atoi(str_ar_select) : opt_ar_select; opt_mbuf = str_mbuf ? strtoul(str_mbuf, NULL, 0) * MiB : opt_mbuf; opt_vmem = str_vmem ? strtoul(str_vmem, NULL, 0) * MiB : opt_vmem; + opt_hostbuf = str_hostbuf ? atoi(str_hostbuf) != 0 : opt_hostbuf; + + // Parse device configuration + const char * str_devices = getenv("GGML_HEXAGON_DEVICES"); + if (!str_devices && str_ndev && str_ndev[0] != '\0') { + GGML_LOG_WARN("DEPRECATED: GGML_HEXAGON_NDEV is deprecated. use GGML_HEXAGON_DEVICES instead\n"); + str_devices = str_ndev; + } + + if (str_devices && str_devices[0] != '\0') { + bool is_single_number = true; + for (int i = 0; str_devices[i] != '\0'; i++) { + if (!isdigit((unsigned char)str_devices[i])) { + is_single_number = false; + break; + } + } + if (is_single_number) { + int n = atoi(str_devices); + if (n < 1) n = 1; + if (n > GGML_HEXAGON_MAX_SESSIONS) n = GGML_HEXAGON_MAX_SESSIONS; + opt_ndev = n; + for (size_t i = 0; i < opt_ndev; i++) { + opt_device_configs[i].physical_idx = 0; + opt_device_configs[i].virtual_idx = (int)i; + opt_device_configs[i].name = "HTP" + std::to_string(i); + } + } else { + std::string s_devices(str_devices); + std::stringstream ss(s_devices); + std::string item; + opt_ndev = 0; + while (std::getline(ss, item, ',')) { + size_t start = item.find_first_not_of(" \t\r\n"); + size_t end = item.find_last_not_of(" \t\r\n"); + if (start == std::string::npos) { + continue; + } + item = item.substr(start, end - start + 1); + + if (item.rfind("HTP", 0) == 0) { + std::string rest = item.substr(3); + size_t colon_pos = rest.find(':'); + int phys = 0; + int virt = 0; + try { + if (colon_pos == std::string::npos) { + phys = std::stoi(rest); + virt = 0; + } else { + phys = std::stoi(rest.substr(0, colon_pos)); + virt = std::stoi(rest.substr(colon_pos + 1)); + } + } catch (...) { + GGML_LOG_WARN("ggml-hex: failed to parse device index in '%s'\n", item.c_str()); + continue; + } - if (opt_ndev > GGML_HEXAGON_MAX_SESSIONS) { - opt_ndev = GGML_HEXAGON_MAX_SESSIONS; + if (opt_ndev < GGML_HEXAGON_MAX_SESSIONS) { + opt_device_configs[opt_ndev].physical_idx = phys; + opt_device_configs[opt_ndev].virtual_idx = virt; + opt_device_configs[opt_ndev].name = colon_pos == std::string::npos + ? "HTP" + std::to_string(phys) + : "HTP" + std::to_string(phys) + ":" + std::to_string(virt); + opt_ndev++; + } else { + GGML_LOG_WARN("ggml-hex: max sessions limit reached (%d), ignoring device %s\n", GGML_HEXAGON_MAX_SESSIONS, item.c_str()); + } + } else { + GGML_LOG_WARN("ggml-hex: invalid device name format '%s', must start with HTP\n", item.c_str()); + } + } + } + } else { + opt_ndev = 1; + opt_device_configs[0].physical_idx = 0; + opt_device_configs[0].virtual_idx = 0; + opt_device_configs[0].name = "HTP0"; } #if defined(__ANDROID__) diff --git a/ggml/src/ggml-hexagon/htp-opnode.h b/ggml/src/ggml-hexagon/htp-opnode.h index b0c859dacf9..741b5e04eb8 100644 --- a/ggml/src/ggml-hexagon/htp-opnode.h +++ b/ggml/src/ggml-hexagon/htp-opnode.h @@ -8,60 +8,107 @@ #include #include #include +#include #include #include "htp-ops.h" #include "htp/matmul-ops.h" #include "htp/flash-attn-ops.h" #include "htp/unary-ops.h" +#include "htp/allreduce-ops.h" struct htp_opnode { - ggml_tensor * node = nullptr; - - std::vector fused; - - htp_op_code opcode = HTP_OP_INVALID; + ggml_tensor * node { nullptr }; + htp_op_code opcode { HTP_OP_INVALID }; + int32_t kernel_params[HTP_OP_MAX_KERN_PARAMS] {0}; + + std::vector fused; + std::vector> dummy; + + std::vector inputs; + std::vector outputs; + std::string name; + + int n_active_src(const ggml_tensor * t) const { + if (!t) return 0; + for (int i = GGML_MAX_SRC - 1; i >= 0; i--) { + if (t->src[i]) { + return i + 1; + } + } + return 0; + } - std::vector extra_dsts; + void init(ggml_tensor * node) { + this->node = node; + if (this->node) { + this->name = ggml_op_desc(this->node); - int32_t kernel_params[HTP_OP_MAX_KERN_PARAMS] = {0}; + // Build inputs (preserving optional nullptrs) + int n_inputs = n_active_src(this->node); + this->inputs.resize(n_inputs, nullptr); + for (int i = 0; i < n_inputs; i++) { + this->inputs[i] = this->node->src[i]; + } - htp_opnode(ggml_tensor * node = nullptr, std::vector fused = {}, htp_op_code opcode = HTP_OP_INVALID, std::vector extra_dsts = {}) - : node(node), fused(std::move(fused)), opcode(opcode), extra_dsts(std::move(extra_dsts)) {} + // Build outputs + this->outputs.push_back(this->dst()); + } + } - ggml_op op() const { - return node->op; + htp_opnode(htp_op_code opcode = HTP_OP_INVALID, ggml_tensor * node = nullptr) : opcode(opcode) { + init(node); } - const ggml_tensor * dst() const { - return fused.empty() ? node : fused.back(); + ggml_op op() const { return node->op; } + const ggml_tensor * src0() const { return node->src[0]; } + const ggml_tensor * src1() const { return node->src[1]; } + const ggml_tensor * dst() const { return outputs.empty() ? node : outputs.back(); } + + ggml_tensor * add_dummy(const ggml_tensor & t) { + dummy.push_back(std::make_shared(t)); + return dummy.back().get(); } void add_fused(ggml_tensor * t, bool extra_dst = false) { fused.push_back(t); + + name += "+"; + name += ggml_op_desc(t); + if (extra_dst) { - extra_dsts.push_back(t); + outputs.push_back(t); + } else { + outputs.clear(); + outputs.push_back(t); } - } - std::vector get_outputs() const { - std::vector res; - if (extra_dsts.empty()) { - res.push_back(dst()); - } else { - res.push_back(node); - for (const auto * x : extra_dsts) { - res.push_back(x); + // Remove the newly fused intermediate output tensor t from inputs (if it was there) + inputs.erase(std::remove(inputs.begin(), inputs.end(), t), inputs.end()); + + // Append new inputs from t, preserving middle nullptrs + int n_inputs = n_active_src(t); + for (int i = 0; i < n_inputs; i++) { + const auto * src = t->src[i]; + if (!src) { + inputs.push_back(nullptr); + } else if (src != node && + std::find(fused.begin(), fused.end(), src) == fused.end() && + std::find(inputs.begin(), inputs.end(), src) == inputs.end()) { + inputs.push_back(src); } } - return res; } - const ggml_tensor * src0() const { - return node->src[0]; + const std::vector & get_inputs() const { + return inputs; } - const ggml_tensor * src1() const { - return node->src[1]; + const std::vector & get_outputs() const { + return outputs; + } + + std::string op_name() const { + return name; } bool is_empty() const { @@ -81,75 +128,6 @@ struct htp_opnode { bool same_input(const htp_opnode& n) const { return n.src1() == this->src1(); } - - std::vector get_inputs() const { - if (fused.empty()) { - int last_non_null = -1; - for (int i = 0; i < GGML_MAX_SRC; i++) { - if (node->src[i]) { - last_non_null = i; - } - } - std::vector inputs(last_non_null + 1, nullptr); - for (int i = 0; i <= last_non_null; i++) { - inputs[i] = node->src[i]; - } - return inputs; - } - - std::vector inputs(GGML_MAX_SRC, nullptr); - std::vector outputs; - outputs.push_back(node); - for (const auto * f : fused) { - outputs.push_back(f); - } - - auto contains = [&](const std::vector & vec, const ggml_tensor * t) { - for (const auto * x : vec) { - if (x == t) return true; - } - return false; - }; - - int count = 0; - auto add_input = [&](const ggml_tensor * t) { - if (t && !contains(outputs, t) && !contains(inputs, t)) { - if (count < (int)inputs.size()) { - inputs[count++] = t; - } else { - inputs.push_back(t); - } - } - }; - - for (int i = 0; i < GGML_MAX_SRC; i++) { - if (node->src[i]) { - add_input(node->src[i]); - } - } - for (const auto * f : fused) { - for (int i = 0; i < GGML_MAX_SRC; i++) { - if (f->src[i]) { - add_input(f->src[i]); - } - } - } - - inputs.resize(count); - return inputs; - } - - std::string op_name() const { - if (fused.empty()) { - return ggml_op_desc(node); - } - std::string name = ggml_op_desc(node); - for (const auto * f : fused) { - name += "+"; - name += ggml_op_desc(f); - } - return name; - } }; struct htp_opformat { @@ -337,8 +315,7 @@ struct htp_opformat { } void format_kernel_params(char * str, size_t max_size, const htp_opnode & node) { if (node.opcode == HTP_OP_MUL_MAT || node.opcode == HTP_OP_MUL_MAT_ID || - node.opcode == HTP_OP_MUL_MAT_QKV || node.opcode == HTP_OP_MUL_MAT_FFN || - node.opcode == HTP_OP_MUL_MAT_ADD) { + node.opcode == HTP_OP_MUL_MAT_NX || node.opcode == HTP_OP_MUL_MAT_ADD) { const auto * kparams = (const struct htp_mm_kernel_params *) node.kernel_params; const char * path = "unknown"; int32_t type = kparams->kernel_type; diff --git a/ggml/src/ggml-hexagon/htp/CMakeLists.txt b/ggml/src/ggml-hexagon/htp/CMakeLists.txt index b00aa2bc94c..77f3ee39dd3 100644 --- a/ggml/src/ggml-hexagon/htp/CMakeLists.txt +++ b/ggml/src/ggml-hexagon/htp/CMakeLists.txt @@ -43,6 +43,7 @@ add_library(${HTP_LIB} SHARED pad-ops.c argsort-ops.c im2col-ops.c + allreduce-ops.c ) target_compile_definitions(${HTP_LIB} PRIVATE diff --git a/ggml/src/ggml-hexagon/htp/act-ops.c b/ggml/src/ggml-hexagon/htp/act-ops.c index 9973c088dda..0a8bf84e382 100644 --- a/ggml/src/ggml-hexagon/htp/act-ops.c +++ b/ggml/src/ggml-hexagon/htp/act-ops.c @@ -183,6 +183,53 @@ static void swiglu_oai_f32(const float * restrict src0, static const float GELU_COEF_A = 0.044715f; static const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f; +static inline HVX_Vector hvx_vec_fast_sigmoid_f32_2it(HVX_Vector v) { + v = Q6_Vqf32_vmpy_VsfVsf(v, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F)); + v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), Q6_V_vsplat_R(FAST_SIGMOID_C3)); + + HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v)); + HVX_Vector x = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int)); + HVX_Vector xx = Q6_Vqf32_vmpy_Vqf32Vqf32(x, x); + + HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx), Q6_V_vsplat_R(FAST_SIGMOID_C2)); + v1 = Q6_Vqf32_vadd_Vqf32Vsf(v1, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F)); + + HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x), Q6_V_vsplat_R(FAST_SIGMOID_C1)); + v2 = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx); + v2 = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x); + + HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1)); + v3 = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24); + + HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1)); + HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4)); + + // Newton-Raphson with 2 iterations + HVX_Vector two_sf = hvx_vec_splat_f32(2.0f); + HVX_Vector i_sf = Q6_Vw_vsub_VwVw(Q6_V_vsplat_R(0x7EEEEBB3), v5); + HVX_Vector r_qf = Q6_Vqf32_vmpy_VsfVsf( + i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v5))))); + r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32( + r_qf, Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v5)))); + HVX_Vector res = Q6_Vsf_equals_Vqf32(r_qf); + + res = Q6_Vqf32_vmpy_VsfVsf(v3, res); + + return Q6_Vsf_equals_Vqf32(res); +} + +static inline HVX_Vector hvx_vec_fast_sigmoid_f32_guard_2it(HVX_Vector v, + HVX_Vector one, + HVX_Vector max_exp, + HVX_Vector min_exp) { + const HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(max_exp, v); + const HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(v, min_exp); + + HVX_Vector out = hvx_vec_fast_sigmoid_f32_2it(v); + out = Q6_V_vmux_QVV(pred_max, out, one); + return Q6_V_vmux_QVV(pred_min, out, Q6_V_vzero()); +} + static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { assert((unsigned long) dst % 128 == 0); assert((unsigned long) src0 % 128 == 0); @@ -200,20 +247,13 @@ static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * rest const HVX_Vector v_coef_a_times_sqrt = hvx_vec_splat_f32(GELU_COEF_A_TIMES_SQRT); const HVX_Vector v_sqrt_2_pi = hvx_vec_splat_f32(SQRT_2_OVER_PI); - const HVX_Vector v_half = hvx_vec_splat_f32(0.5f); const HVX_Vector v_one = hvx_vec_splat_f32(1.0f); - const HVX_Vector v_two = hvx_vec_splat_f32(2.0f); - - // Hoisted fast sigmoid / inverse constants to avoid loop-internal overhead - const HVX_Vector v_log2f = Q6_V_vsplat_R(FAST_SIGMOID_LOG2F); - const HVX_Vector v_c1 = Q6_V_vsplat_R(FAST_SIGMOID_C1); - const HVX_Vector v_c2 = Q6_V_vsplat_R(FAST_SIGMOID_C2); - const HVX_Vector v_inv_aprox = Q6_V_vsplat_R(0x7EEEEBB3); const HVX_Vector v_max_exp = hvx_vec_splat_f32(87.0f); const HVX_Vector v_min_exp = hvx_vec_splat_f32(-87.0f); uint32_t i = 0; + _Pragma("unroll(4)") for (; i < nvec; i++) { HVX_Vector x = vsrc0[i]; HVX_Vector g = vsrc1[i]; @@ -223,56 +263,13 @@ static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * rest coef = hvx_vec_add_f32_f32(coef, v_sqrt_2_pi); HVX_Vector inner = hvx_vec_mul_f32_f32(x, coef); - // y2 = 2 * inner - HVX_Vector y2 = hvx_vec_mul_f32_f32(inner, v_two); - - // Sigmoid guard check predicates - HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(v_max_exp, y2); - HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(y2, v_min_exp); - - // Fast sigmoid approximation - HVX_Vector v = Q6_Vqf32_vmpy_VsfVsf(y2, v_log2f); - v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), v_half); - - HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v)); - HVX_Vector x_sig = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int)); - HVX_Vector xx_sig = Q6_Vqf32_vmpy_Vqf32Vqf32(x_sig, x_sig); - - HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx_sig), v_c2); - v1 = Q6_Vqf32_vadd_Vqf32Vsf(v1, v_log2f); + // y2 = 2 * inner = inner + inner + HVX_Vector y2 = hvx_vec_add_f32_f32(inner, inner); - HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x_sig), v_c1); - v2 = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx_sig); - v2 = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x_sig); - - HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1)); - v3 = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24); - - HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1)); - HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4)); - - // Fast division (Newton-Raphson with 2 iterations) - HVX_Vector i_sf = Q6_Vw_vsub_VwVw(v_inv_aprox, v5); - HVX_Vector r_qf = Q6_Vqf32_vmpy_VsfVsf( - i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v_two, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v5))))); - r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32( - r_qf, Q6_Vqf32_vsub_VsfVsf(v_two, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v5)))); - HVX_Vector res_inv = Q6_Vsf_equals_Vqf32(r_qf); - - HVX_Vector sig2y = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v3, res_inv)); - - // Sigmoid guards - sig2y = Q6_V_vmux_QVV(pred_max, sig2y, v_one); - sig2y = Q6_V_vmux_QVV(pred_min, sig2y, Q6_V_vzero()); - - // tanh(inner) = 2 * sigmoid(2 * inner) - 1 - HVX_Vector tanh_val = hvx_vec_mul_f32_f32(sig2y, v_two); - tanh_val = hvx_vec_sub_f32_f32(tanh_val, v_one); - - HVX_Vector tanh_plus_one = hvx_vec_add_f32_f32(tanh_val, v_one); - HVX_Vector half_x = hvx_vec_mul_f32_f32(x, v_half); - HVX_Vector gelu_x = hvx_vec_mul_f32_f32(half_x, tanh_plus_one); + // Fast sigmoid approximation (2 iterations) + HVX_Vector sig2y = hvx_vec_fast_sigmoid_f32_guard_2it(y2, v_one, v_max_exp, v_min_exp); + HVX_Vector gelu_x = hvx_vec_mul_f32_f32(x, sig2y); vdst[i] = hvx_vec_mul_f32_f32(gelu_x, g); } @@ -285,50 +282,11 @@ static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * rest coef = hvx_vec_add_f32_f32(coef, v_sqrt_2_pi); HVX_Vector inner = hvx_vec_mul_f32_f32(x, coef); - HVX_Vector y2 = hvx_vec_mul_f32_f32(inner, v_two); - - HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(v_max_exp, y2); - HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(y2, v_min_exp); - - HVX_Vector v = Q6_Vqf32_vmpy_VsfVsf(y2, v_log2f); - v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), v_half); - - HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v)); - HVX_Vector x_sig = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int)); - HVX_Vector xx_sig = Q6_Vqf32_vmpy_Vqf32Vqf32(x_sig, x_sig); - - HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx_sig), v_c2); - v1 = Q6_Vqf32_vadd_Vqf32Vsf(v1, v_log2f); - - HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x_sig), v_c1); - v2 = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx_sig); - v2 = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x_sig); - - HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1)); - v3 = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24); - - HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1)); - HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4)); - - HVX_Vector i_sf = Q6_Vw_vsub_VwVw(v_inv_aprox, v5); - HVX_Vector r_qf = Q6_Vqf32_vmpy_VsfVsf( - i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v_two, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v5))))); - r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32( - r_qf, Q6_Vqf32_vsub_VsfVsf(v_two, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v5)))); - HVX_Vector res_inv = Q6_Vsf_equals_Vqf32(r_qf); - - HVX_Vector sig2y = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v3, res_inv)); - - sig2y = Q6_V_vmux_QVV(pred_max, sig2y, v_one); - sig2y = Q6_V_vmux_QVV(pred_min, sig2y, Q6_V_vzero()); - - HVX_Vector tanh_val = hvx_vec_mul_f32_f32(sig2y, v_two); - tanh_val = hvx_vec_sub_f32_f32(tanh_val, v_one); + HVX_Vector y2 = hvx_vec_add_f32_f32(inner, inner); - HVX_Vector tanh_plus_one = hvx_vec_add_f32_f32(tanh_val, v_one); - HVX_Vector half_x = hvx_vec_mul_f32_f32(x, v_half); - HVX_Vector gelu_x = hvx_vec_mul_f32_f32(half_x, tanh_plus_one); + HVX_Vector sig2y = hvx_vec_fast_sigmoid_f32_guard_2it(y2, v_one, v_max_exp, v_min_exp); + HVX_Vector gelu_x = hvx_vec_mul_f32_f32(x, sig2y); HVX_Vector res = hvx_vec_mul_f32_f32(gelu_x, g); hvx_vec_store_a((void *) &vdst[i], nloe * sizeof(float), res); } diff --git a/ggml/src/ggml-hexagon/htp/allreduce-ops.c b/ggml/src/ggml-hexagon/htp/allreduce-ops.c new file mode 100644 index 00000000000..d35f685a6dc --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/allreduce-ops.c @@ -0,0 +1,398 @@ +#pragma clang diagnostic ignored "-Wunused-variable" +#pragma clang diagnostic ignored "-Wunused-function" +#pragma clang diagnostic ignored "-Wunused-but-set-variable" + +#include +#include +#include +#include +#include + +#define GGML_COMMON_DECL_C +#include "ggml-common.h" +#include "htp-ctx.h" +#include "htp-ops.h" +#include "hvx-utils.h" +#include "htp-tensor.h" +#include "hex-dma.h" +#include "hex-profile.h" +#include "allreduce-ops.h" + +struct htp_allreduce_context { + struct htp_ops_context * octx; + uint32_t n_ranks; + uint32_t n_dsts; + uint32_t nelem; + uint32_t ne0; + uint32_t ne1; + uint32_t row_size_aligned; + uint32_t rank_elem_start; + uint32_t rank_nelem; + uint32_t elems_per_thread; + uint32_t block_elems; + uint32_t vtcm_size_per_thread; + bool is_row_bcast; + uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS]; + uint8_t * dst_spad_base; + uint8_t * res_spad_base; +}; + +#define DEFINE_ALLREDUCE_THREAD_DMA_1D(SUFFIX, TYPE, HVX_ADD_FN, HAS_ADD) \ +static void allreduce_thread_dma_1d_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ + struct htp_allreduce_context * actx = (struct htp_allreduce_context *) data; \ + struct htp_ops_context * octx = actx->octx; \ + \ + const uint32_t n_ranks = actx->n_ranks; \ + const uint32_t n_dsts = actx->n_dsts; \ + const uint32_t block_elems = actx->block_elems; \ + \ + const uint32_t dr = actx->elems_per_thread; \ + const uint32_t ir0 = actx->rank_elem_start + dr * ith; \ + const uint32_t ir1 = MIN(ir0 + dr, actx->rank_elem_start + actx->rank_nelem); \ + if (ir0 >= ir1) return; \ + \ + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ + dma_queue * q = octx->ctx->dma[ith]; \ + \ + uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS]; \ + for (uint32_t s = 0; s < n_ranks; s++) { \ + src_spad_base[s] = actx->src_spad_base[s] + (ith * actx->vtcm_size_per_thread); \ + } \ + uint8_t * dst_spad_base = actx->dst_spad_base + (ith * actx->vtcm_size_per_thread); \ + uint8_t * res_spad_base = HAS_ADD ? (actx->res_spad_base + (ith * actx->vtcm_size_per_thread)) : NULL; \ + \ + const size_t spad_half = actx->vtcm_size_per_thread / 2; \ + uint32_t ir_prefetch = ir0; \ + int spad_idx = 0; \ + \ + for (int k = 0; k < 2 && ir_prefetch < ir1; k++) { \ + uint32_t cur_elems = MIN(block_elems, ir1 - ir_prefetch); \ + size_t cur_bytes = cur_elems * sizeof(TYPE); \ + uint8_t * d_spad = dst_spad_base + spad_idx * spad_half; \ + for (uint32_t d = 0; d < n_dsts; d++) { \ + uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + ir_prefetch * sizeof(TYPE); \ + dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 0); \ + } \ + for (uint32_t s = 0; s < n_ranks; s++) { \ + uint8_t * s_spad = src_spad_base[s] + spad_idx * spad_half; \ + const uint8_t * s_ddr = (const uint8_t *) octx->src[s]->data + ir_prefetch * sizeof(TYPE); \ + dma_queue_push(q, dma_make_ptr(s_spad, s_ddr), cur_bytes, cur_bytes, cur_bytes, 1); \ + } \ + if (HAS_ADD) { \ + uint8_t * r_spad = res_spad_base + spad_idx * spad_half; \ + const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE); \ + dma_queue_push(q, dma_make_ptr(r_spad, r_ddr), cur_bytes, cur_bytes, cur_bytes, 1); \ + } \ + ir_prefetch += cur_elems; \ + spad_idx ^= 1; \ + } \ + \ + for (uint32_t ir = ir0; ir < ir1; ) { \ + uint32_t cur_elems = MIN(block_elems, ir1 - ir); \ + size_t cur_bytes = cur_elems * sizeof(TYPE); \ + uint8_t * d_spad = NULL; \ + for (uint32_t d = 0; d < n_dsts; d++) { \ + d_spad = (uint8_t *) dma_queue_pop(q).src; \ + } \ + uint8_t * s_spad[HTP_ALLREDUCE_MAX_RANKS]; \ + for (uint32_t s = 0; s < n_ranks; s++) { \ + s_spad[s] = (uint8_t *) dma_queue_pop(q).dst; \ + } \ + uint8_t * r_spad = HAS_ADD ? (uint8_t *) dma_queue_pop(q).dst : NULL; \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); \ + HVX_ADD_FN(d_spad, s_spad[0], s_spad[1], cur_elems); \ + for (uint32_t s = 2; s < n_ranks; s++) { \ + HVX_ADD_FN(d_spad, d_spad, s_spad[s], cur_elems); \ + } \ + if (HAS_ADD) { \ + HVX_ADD_FN(d_spad, d_spad, r_spad, cur_elems); \ + } \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); \ + for (uint32_t d = 0; d < n_dsts; d++) { \ + uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + ir * sizeof(TYPE); \ + dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 1); \ + } \ + if (ir_prefetch < ir1) { \ + uint32_t next_elems = MIN(block_elems, ir1 - ir_prefetch); \ + size_t next_bytes = next_elems * sizeof(TYPE); \ + for (uint32_t s = 0; s < n_ranks; s++) { \ + const uint8_t * s_next = (const uint8_t *) octx->src[s]->data + ir_prefetch * sizeof(TYPE); \ + dma_queue_push(q, dma_make_ptr(s_spad[s], s_next), next_bytes, next_bytes, next_bytes, 1); \ + } \ + if (HAS_ADD) { \ + const uint8_t * r_next = (const uint8_t *) octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE); \ + dma_queue_push(q, dma_make_ptr(r_spad, r_next), next_bytes, next_bytes, next_bytes, 1); \ + } \ + ir_prefetch += next_elems; \ + } \ + ir += cur_elems; \ + } \ + dma_queue_flush(q); \ +} + +DEFINE_ALLREDUCE_THREAD_DMA_1D(f16, __fp16, hvx_add_f16_aaa, 0) +DEFINE_ALLREDUCE_THREAD_DMA_1D(f32, float, hvx_add_f32_aaa, 0) +DEFINE_ALLREDUCE_THREAD_DMA_1D(add_f16, __fp16, hvx_add_f16_aaa, 1) +DEFINE_ALLREDUCE_THREAD_DMA_1D(add_f32, float, hvx_add_f32_aaa, 1) + +#define DEFINE_ALLREDUCE_THREAD_DMA_2D(SUFFIX, TYPE, HVX_ADD_FN, HAS_ADD, IS_ROW_BCAST) \ +static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ + struct htp_allreduce_context * actx = (struct htp_allreduce_context *) data; \ + struct htp_ops_context * octx = actx->octx; \ + \ + const uint32_t n_ranks = actx->n_ranks; \ + const uint32_t n_dsts = actx->n_dsts; \ + const uint32_t ne0 = actx->ne0; \ + const uint32_t block_rows = actx->block_elems; \ + const uint32_t row_size_aligned = actx->row_size_aligned; \ + const uint32_t row_bytes = ne0 * sizeof(TYPE); \ + \ + const uint32_t dr = actx->elems_per_thread; \ + const uint32_t r0 = actx->rank_elem_start + dr * ith; \ + const uint32_t r1 = MIN(r0 + dr, actx->rank_elem_start + actx->rank_nelem); \ + if (r0 >= r1) return; \ + \ + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ + dma_queue * q = octx->ctx->dma[ith]; \ + \ + uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS]; \ + for (uint32_t s = 0; s < n_ranks; s++) { \ + src_spad_base[s] = actx->src_spad_base[s] + (ith * actx->vtcm_size_per_thread); \ + } \ + uint8_t * dst_spad_base = actx->dst_spad_base + (ith * actx->vtcm_size_per_thread); \ + uint8_t * res_spad_base = HAS_ADD ? (IS_ROW_BCAST ? actx->res_spad_base : (actx->res_spad_base + (ith * actx->vtcm_size_per_thread))) : NULL; \ + \ + const size_t spad_half = actx->vtcm_size_per_thread / 2; \ + uint32_t r_prefetch = r0; \ + int spad_idx = 0; \ + \ + for (int k = 0; k < 2 && r_prefetch < r1; k++) { \ + uint32_t cur_rows = MIN(block_rows, r1 - r_prefetch); \ + uint8_t * d_spad = dst_spad_base + spad_idx * spad_half; \ + for (uint32_t d = 0; d < n_dsts; d++) { \ + uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + r_prefetch * octx->dsts[d]->nb[1]; \ + dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, 0); \ + } \ + for (uint32_t s = 0; s < n_ranks; s++) { \ + uint8_t * s_spad = src_spad_base[s] + spad_idx * spad_half; \ + const uint8_t * s_ddr = (const uint8_t *) octx->src[s]->data + r_prefetch * octx->src[s]->nb[1]; \ + dma_queue_push(q, dma_make_ptr(s_spad, s_ddr), row_size_aligned, octx->src[s]->nb[1], row_bytes, cur_rows); \ + } \ + if (HAS_ADD && !IS_ROW_BCAST) { \ + uint8_t * r_spad = res_spad_base + spad_idx * spad_half; \ + const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data + r_prefetch * octx->src[2 * n_ranks]->nb[1]; \ + dma_queue_push(q, dma_make_ptr(r_spad, r_ddr), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, cur_rows); \ + } \ + r_prefetch += cur_rows; \ + spad_idx ^= 1; \ + } \ + \ + for (uint32_t r = r0; r < r1; ) { \ + uint32_t cur_rows = MIN(block_rows, r1 - r); \ + uint8_t * d_spad = NULL; \ + for (uint32_t d = 0; d < n_dsts; d++) { \ + d_spad = (uint8_t *) dma_queue_pop(q).src; \ + } \ + uint8_t * s_spad[HTP_ALLREDUCE_MAX_RANKS]; \ + for (uint32_t s = 0; s < n_ranks; s++) { \ + s_spad[s] = (uint8_t *) dma_queue_pop(q).dst; \ + } \ + uint8_t * r_spad = (HAS_ADD && !IS_ROW_BCAST) ? (uint8_t *) dma_queue_pop(q).dst : NULL; \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r); \ + for (uint32_t row = 0; row < cur_rows; row++) { \ + uint8_t * d_row = d_spad + row * row_size_aligned; \ + const uint8_t * s0_row = s_spad[0] + row * row_size_aligned; \ + const uint8_t * s1_row = s_spad[1] + row * row_size_aligned; \ + HVX_ADD_FN(d_row, s0_row, s1_row, ne0); \ + for (uint32_t s = 2; s < n_ranks; s++) { \ + const uint8_t * ss_row = s_spad[s] + row * row_size_aligned; \ + HVX_ADD_FN(d_row, d_row, ss_row, ne0); \ + } \ + if (HAS_ADD) { \ + const uint8_t * res_row = IS_ROW_BCAST ? res_spad_base : (r_spad + row * row_size_aligned); \ + HVX_ADD_FN(d_row, d_row, res_row, ne0); \ + } \ + } \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r); \ + for (uint32_t d = 0; d < n_dsts; d++) { \ + uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + r * octx->dsts[d]->nb[1]; \ + dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, cur_rows); \ + } \ + if (r_prefetch < r1) { \ + uint32_t next_rows = MIN(block_rows, r1 - r_prefetch); \ + for (uint32_t s = 0; s < n_ranks; s++) { \ + const uint8_t * s_next = (const uint8_t *) octx->src[s]->data + r_prefetch * octx->src[s]->nb[1]; \ + dma_queue_push(q, dma_make_ptr(s_spad[s], s_next), row_size_aligned, octx->src[s]->nb[1], row_bytes, next_rows); \ + } \ + if (HAS_ADD && !IS_ROW_BCAST) { \ + const uint8_t * r_next = (const uint8_t *) octx->src[2 * n_ranks]->data + r_prefetch * octx->src[2 * n_ranks]->nb[1]; \ + dma_queue_push(q, dma_make_ptr(r_spad, r_next), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, next_rows); \ + } \ + r_prefetch += next_rows; \ + } \ + r += cur_rows; \ + } \ + dma_queue_flush(q); \ +} + +DEFINE_ALLREDUCE_THREAD_DMA_2D(f16, __fp16, hvx_add_f16_aaa, 0, 0) +DEFINE_ALLREDUCE_THREAD_DMA_2D(f32, float, hvx_add_f32_aaa, 0, 0) +DEFINE_ALLREDUCE_THREAD_DMA_2D(add_f16, __fp16, hvx_add_f16_aaa, 1, 0) +DEFINE_ALLREDUCE_THREAD_DMA_2D(add_f32, float, hvx_add_f32_aaa, 1, 0) +DEFINE_ALLREDUCE_THREAD_DMA_2D(add_bcast_f16, __fp16, hvx_add_f16_aaa, 1, 1) +DEFINE_ALLREDUCE_THREAD_DMA_2D(add_bcast_f32, float, hvx_add_f32_aaa, 1, 1) + +int op_allreduce(struct htp_ops_context * octx) { + const struct htp_allreduce_kernel_params * kparams = (const struct htp_allreduce_kernel_params *) octx->kernel_params; + const struct htp_tensor * dst = octx->dst; + + const uint32_t rank = (uint32_t) kparams->rank; + const uint32_t n_ranks = (uint32_t) kparams->n_ranks; + + if (n_ranks < 2 || n_ranks > HTP_ALLREDUCE_MAX_RANKS || rank >= n_ranks) { + return HTP_STATUS_INVAL_PARAMS; + } + + if (dst->type != HTP_TYPE_F16 && dst->type != HTP_TYPE_F32) { + return HTP_STATUS_NO_SUPPORT; + } + + const uint32_t nelem = dst->ne[0] * dst->ne[1] * dst->ne[2] * dst->ne[3]; + const uint32_t fence_seq_entry = (uint32_t) octx->op_params[0]; + const uint32_t fence_seq_exit = (uint32_t) octx->op_params[1]; + + // 1. Entry Barrier: Synchronize all ranks before reading + struct htp_thread_trace * tr0 = &octx->ctx->trace[0]; + htp_trace_event_start(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry); + + const struct htp_tensor * my_sync = octx->src[n_ranks + rank]; + atomic_uint * my_fence = (atomic_uint *) my_sync->data; + + atomic_store(&my_fence[0], fence_seq_entry); + asm volatile ("syncht" : : : "memory"); + Q6_dccleaninva_A((void *) my_fence); + + for (uint32_t j = 0; j < n_ranks; j++) { + if (j == rank) continue; + const struct htp_tensor * peer_sync = octx->src[n_ranks + j]; + atomic_uint * peer_fence = (atomic_uint *) peer_sync->data; + uint64_t spins = 0; + while (1) { + Q6_dccleaninva_A((void *) peer_fence); + uint32_t val = atomic_load(&peer_fence[0]); + if (val == fence_seq_entry || val == fence_seq_exit) { + break; + } + if (++spins > HTP_FENCE_TIMEOUT) { + FARF(ERROR, "ggml-hex: allreduce entry fence-wait TIMEOUT: rank %u waiting on %u (fence %p seq %u)\n", rank, j, peer_fence, fence_seq_entry); + return HTP_STATUS_INTERNAL_ERR; + } + hex_pause(); + } + } + asm volatile ("syncht" : : : "memory"); + + htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry); + + // 2. Multi-threaded Reduction across assigned rank chunk + if (nelem > 0) { + const uint32_t n_threads = (uint32_t) kparams->n_threads; + const uint32_t block_elems = (uint32_t) kparams->block_elems; + const uint32_t elems_per_thread = (uint32_t) kparams->elems_per_thread; + const uint32_t vtcm_size_per_thread = (uint32_t) kparams->vtcm_size_per_thread; + + const bool has_add = (octx->op == HTP_OP_ALLREDUCE_ADD); + + struct htp_allreduce_context actx; + actx.octx = octx; + actx.n_ranks = n_ranks; + actx.n_dsts = (uint32_t) kparams->n_dsts ? (uint32_t) kparams->n_dsts : n_ranks; + actx.nelem = nelem; + actx.ne0 = (uint32_t) kparams->ne0; + actx.ne1 = (uint32_t) kparams->ne1; + actx.row_size_aligned = (uint32_t) kparams->row_size_aligned; + actx.rank_elem_start = (uint32_t) kparams->rank_elem_start; + actx.rank_nelem = (uint32_t) kparams->rank_nelem; + actx.elems_per_thread = elems_per_thread; + actx.block_elems = block_elems; + actx.vtcm_size_per_thread = vtcm_size_per_thread; + actx.is_row_bcast = (kparams->is_row_bcast != 0); + + work_queue_func_t reduce_fun = NULL; + switch (kparams->kernel_type) { + case HTP_ALLREDUCE_KERNEL_DMA_1D: + if (has_add) { + reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_1d_add_f16 : allreduce_thread_dma_1d_add_f32; + } else { + reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_1d_f16 : allreduce_thread_dma_1d_f32; + } + break; + case HTP_ALLREDUCE_KERNEL_DMA_2D: + if (has_add) { + if (kparams->is_row_bcast) { + reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_2d_add_bcast_f16 : allreduce_thread_dma_2d_add_bcast_f32; + } else { + reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_2d_add_f16 : allreduce_thread_dma_2d_add_f32; + } + } else { + reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_2d_f16 : allreduce_thread_dma_2d_f32; + } + break; + default: + return HTP_STATUS_NO_SUPPORT; + } + + uint8_t * vtcm_ptr = (uint8_t *) octx->ctx->vtcm_base; + for (uint32_t s = 0; s < n_ranks; s++) { + actx.src_spad_base[s] = vtcm_ptr; + vtcm_ptr += n_threads * vtcm_size_per_thread; + } + actx.dst_spad_base = vtcm_ptr; + vtcm_ptr += n_threads * vtcm_size_per_thread; + if (has_add) { + actx.res_spad_base = vtcm_ptr; + vtcm_ptr += (actx.is_row_bcast ? 1 : n_threads) * vtcm_size_per_thread; + } + + if (has_add && actx.is_row_bcast) { + const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data; + const uint32_t row_bytes = actx.ne0 * (dst->type == HTP_TYPE_F16 ? sizeof(__fp16) : sizeof(float)); + dma_queue * q = octx->ctx->dma[0]; + dma_queue_push(q, dma_make_ptr(actx.res_spad_base, r_ddr), actx.row_size_aligned, 0, row_bytes, 1); + dma_queue_pop(q); + } + + work_queue_run(octx->ctx->work_queue, reduce_fun, &actx, n_threads); + } + + // 4. Exit Barrier: Synchronize all ranks after writing + htp_trace_event_start(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit); + + atomic_store(&my_fence[0], fence_seq_exit); + asm volatile ("syncht" : : : "memory"); + Q6_dccleaninva_A((void *) my_fence); + + for (uint32_t j = 0; j < n_ranks; j++) { + if (j == rank) continue; + const struct htp_tensor * peer_sync = octx->src[n_ranks + j]; + atomic_uint * peer_fence = (atomic_uint *) peer_sync->data; + uint64_t spins = 0; + while (1) { + Q6_dccleaninva_A((void *) peer_fence); + uint32_t val = atomic_load(&peer_fence[0]); + if (val == fence_seq_exit) { + break; + } + if (++spins > HTP_FENCE_TIMEOUT) { + FARF(ERROR, "ggml-hex: allreduce exit fence-wait TIMEOUT: rank %u waiting on %u (fence %p seq %u)\n", rank, j, peer_fence, fence_seq_exit); + return HTP_STATUS_INTERNAL_ERR; + } + hex_pause(); + } + } + asm volatile ("syncht" : : : "memory"); + + htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit); + + return HTP_STATUS_OK; +} diff --git a/ggml/src/ggml-hexagon/htp/allreduce-ops.h b/ggml/src/ggml-hexagon/htp/allreduce-ops.h new file mode 100644 index 00000000000..de447d87e91 --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/allreduce-ops.h @@ -0,0 +1,40 @@ +#ifndef ALLREDUCE_OPS_H +#define ALLREDUCE_OPS_H + +#include + +#define HTP_ALLREDUCE_MAX_RANKS 4 + +#ifdef __cplusplus +extern "C" { +#endif + +enum htp_allreduce_kernel_type { + HTP_ALLREDUCE_KERNEL_UNSUPPORTED = 0, + HTP_ALLREDUCE_KERNEL_DMA_1D, + HTP_ALLREDUCE_KERNEL_DMA_2D, +}; + +struct htp_allreduce_kernel_params { + int32_t rank; + int32_t n_ranks; + int32_t n_threads; + int32_t block_elems; // 1D: block_elems, 2D: block_rows + int32_t elems_per_thread; // 1D: nelem_per_thread, 2D: nrows_per_thread + int32_t vtcm_size_per_thread; + int32_t vtcm_size; + int32_t kernel_type; + int32_t ne0; + int32_t ne1; + int32_t row_size_aligned; + int32_t rank_elem_start; + int32_t rank_nelem; + int32_t n_dsts; + int32_t is_row_bcast; +}; + +#ifdef __cplusplus +} +#endif + +#endif /* ALLREDUCE_OPS_H */ diff --git a/ggml/src/ggml-hexagon/htp/cpy-ops.c b/ggml/src/ggml-hexagon/htp/cpy-ops.c index ae507effa51..15bc8dc244f 100644 --- a/ggml/src/ggml-hexagon/htp/cpy-ops.c +++ b/ggml/src/ggml-hexagon/htp/cpy-ops.c @@ -4,6 +4,7 @@ #include #include +#include #include #include @@ -14,6 +15,7 @@ #include "htp-ops.h" #include "htp-ops.h" #include "hvx-utils.h" +#include "htp-tensor.h" struct htp_copy_context { struct htp_ops_context * octx; @@ -78,7 +80,7 @@ static void cpy_thread_##NAME##_sameshape(unsigned int nth, unsigned int ith, vo } \ } -DEFINE_CPY_SAMESHAPE(f32, float, 4) +DEFINE_CPY_SAMESHAPE(f32, float, 4) DEFINE_CPY_SAMESHAPE(f16, __fp16, 2) #define DEFINE_CPY_RESHAPE(NAME, ELEM_TYPE, ELEM_SIZE) \ @@ -179,7 +181,7 @@ static void cpy_thread_##NAME##_reshape(unsigned int nth, unsigned int ith, void } \ } -DEFINE_CPY_RESHAPE(f32, float, 4) +DEFINE_CPY_RESHAPE(f32, float, 4) DEFINE_CPY_RESHAPE(f16, __fp16, 2) static void cpy_thread_f16_f32_sameshape(unsigned int nth, unsigned int ith, void * data) { @@ -232,6 +234,41 @@ static void cpy_thread_f32_f16_sameshape(unsigned int nth, unsigned int ith, voi } } +static inline void cpy_dma_sametype_sameshape( + struct htp_ops_context * octx, + const struct htp_tensor * dst, + const struct htp_tensor * src0, + uint32_t elem_size, + uint32_t ne00, uint32_t ne01, uint32_t ne02, uint32_t ne03, + uint32_t nb01, uint32_t nb02, uint32_t nb03, + uint32_t nb1, uint32_t nb2, uint32_t nb3 +) { + const bool contiguous_outer = + (ne02 == 1 || (nb02 == ne01 * nb01 && nb2 == ne01 * nb1)) && + (ne03 == 1 || (nb03 == ne02 * nb02 && nb3 == ne02 * nb2)); + + dma_queue * q = octx->ctx->dma[0]; + + if (contiguous_outer) { + dma_queue_push(q, dma_make_ptr((void *) dst->data, (const void *) src0->data), nb1, nb01, ne00 * elem_size, ne01 * ne02 * ne03); + dma_queue_pop(q); + return; + } + + for (uint32_t i03 = 0; i03 < ne03; i03++) { + for (uint32_t i02 = 0; i02 < ne02; i02++) { + uint8_t* dst_ptr = (uint8_t*) dst->data + i02*nb2 + i03*nb3; + uint8_t* src0_ptr = (uint8_t*) src0->data + i02*nb02 + i03*nb03; + if (!dma_queue_push(q, dma_make_ptr(dst_ptr, src0_ptr), nb1, nb01, ne00 * elem_size, ne01)) { + dma_queue_flush(q); + dma_queue_push(q, dma_make_ptr(dst_ptr, src0_ptr), nb1, nb01, ne00 * elem_size, ne01); + } + } + } + + dma_queue_flush(q); +} + int op_cpy(struct htp_ops_context * octx) { cpy_preamble; @@ -264,14 +301,11 @@ int op_cpy(struct htp_ops_context * octx) { ct.src0_nrows_per_thread = (nr + n_threads - 1) / n_threads; - worker_callback_t copy_fun; + worker_callback_t copy_fun = NULL; + bool use_dma = false; if (sametype && sameshape) { - if (src0->type == HTP_TYPE_F32) { - copy_fun = cpy_thread_f32_sameshape; - } else { - copy_fun = cpy_thread_f16_sameshape; - } + use_dma = true; } else if (sameshape) { /**/ if (dst->type == HTP_TYPE_F16 && src0->type == HTP_TYPE_F32) copy_fun = cpy_thread_f16_f32_sameshape; @@ -289,7 +323,28 @@ int op_cpy(struct htp_ops_context * octx) { return HTP_STATUS_NO_SUPPORT; } - worker_pool_run_func(octx->ctx->worker_pool, copy_fun, &ct, n_threads); + if (use_dma) { + cpy_dma_sametype_sameshape(octx, dst, src0, ct.src0_type_size, ne00, ne01, ne02, ne03, nb01, nb02, nb03, nb1, nb2, nb3); + } else { + worker_pool_run_func(octx->ctx->worker_pool, copy_fun, &ct, n_threads); + } + + const struct htp_tensor *sync = octx->src[1]; + if (sync) { + if (!use_dma) { + // htp_tensor_flush_all(octx->ctx, octx->dsts, 1); + qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE); + } + + atomic_uint * sync_fence = (atomic_uint *) sync->data; + const uint32_t seq = (uint32_t) octx->op_params[0]; + + atomic_store(&sync_fence[0], seq); + asm volatile ("syncht" : : : "memory"); + Q6_dccleaninva_A((void *) sync_fence); + + FARF(HIGH, "ggml-hex: sync-release : fence %p seq %u\n", sync_fence, seq); + } return HTP_STATUS_OK; } diff --git a/ggml/src/ggml-hexagon/htp/dma-queue.h b/ggml/src/ggml-hexagon/htp/dma-queue.h index 264284bda82..190ca3a9b9e 100644 --- a/ggml/src/ggml-hexagon/htp/dma-queue.h +++ b/ggml/src/ggml-hexagon/htp/dma-queue.h @@ -244,17 +244,18 @@ static inline dma_ptr dma_queue_pop(dma_queue * q) { return dptr; } - dma_descriptor_2d * desc = &r->desc[r->pop_idx]; + dptr = r->dptr[r->pop_idx]; + + volatile dma_descriptor_2d * desc = &r->desc[r->pop_idx]; // Wait for desc to complete if (!desc->done) { + // FARF(ALWAYS, "dma-poll: idx %u dst %p src %p", r->pop_idx, dptr.dst, dptr.src); while (!desc->done) { dmpoll(); } } - dptr = r->dptr[r->pop_idx]; - htp_trace_event_stop(r->trace, HTP_TRACE_EVT_DMA, r->pop_idx); r->pop_idx = (r->pop_idx + 1) & r->idx_mask; diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index 81765629046..c76b4d3a3ac 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -30,6 +30,8 @@ #include "ggml-common.h" #include "htp-ctx.h" #include "htp-ops.h" +#include "htp-tensor.h" +#include "hvx-quant.h" #include "flash-attn-ops.h" #include "hvx-fa-kernels.h" @@ -85,12 +87,17 @@ struct htp_fa_context { uint8_t * spad_m; uint8_t * spad_a; + const struct htp_tensor * k; + const struct htp_tensor * v; + uint64_t t_start; }; struct hmx_fa_context { const struct htp_ops_context * octx; const struct htp_tensor * sinks; // attention sinks (src[4]), NULL if absent + const struct htp_tensor * k; + const struct htp_tensor * v; bool pipeline; // true when n_kv_blocks >= FA_MIN_KV_BLOCKS && n_threads >= 2 uint32_t n_threads; @@ -214,8 +221,8 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * const uint32_t DV = nev0; const size_t size_q_row = DK * ((q->type == HTP_TYPE_F32) ? 4 : 2); - const size_t size_k_row = DK * sizeof(__fp16); - const size_t size_v_row = DV * sizeof(__fp16); + const size_t size_k_row = htp_tensor_get_row_size(k->type, DK); + const size_t size_v_row = htp_tensor_get_row_size(v->type, DV); // Scratchpad buffers for Q, K, V, Mask, and VKQ32 accumulator uint8_t * spad_q = factx->spad_q + factx->size_q_block * ith; @@ -364,6 +371,23 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * uint8_t * v_base = dma_queue_pop(dma).dst; // V __fp16 * m_base = mask ? dma_queue_pop(dma).dst : NULL; // M + if (factx->k->type == HTP_TYPE_Q8_0) { + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, ir); + for (uint32_t r = 0; r < current_block_size; ++r) { + __fp16 * row_k = (__fp16 *)(k_base + r * factx->size_k_row_padded); + hvx_dequantize_row_q8_0_f16(row_k, row_k, DK); + } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, ir); + } + if (factx->v->type == HTP_TYPE_Q8_0) { + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, ir); + for (uint32_t r = 0; r < current_block_size; ++r) { + __fp16 * row_v = (__fp16 *)(v_base + r * factx->size_v_row_padded); + hvx_dequantize_row_q8_0_f16(row_v, row_v, DV); + } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, ir); + } + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_QK, ir); // Inner loop processing the block from VTCM @@ -625,6 +649,12 @@ static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data) struct htp_thread_trace * tr = &factx->octx->ctx->trace[i]; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start)); + if (factx->k->type == HTP_TYPE_Q8_0) { + for (uint32_t r = start; r < end; ++r) { + __fp16 * row_k = (__fp16 *)((char *)args->curr_k + r * args->src_stride * sizeof(__fp16)); + hvx_dequantize_row_q8_0_f16(row_k, row_k, factx->DK); + } + } hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles[args->buf_idx], (const __fp16 *) args->curr_k, total_rows, factx->DK, args->src_stride, start, end); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start)); @@ -673,6 +703,12 @@ static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data) struct htp_thread_trace * tr = &factx->octx->ctx->trace[i]; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, (uint16_t) (args->kv_start + start)); + if (factx->v->type == HTP_TYPE_Q8_0) { + for (uint32_t r = start; r < end; ++r) { + __fp16 * row_v = (__fp16 *)((char *)args->v_src + r * args->src_stride * sizeof(__fp16)); + hvx_dequantize_row_q8_0_f16(row_v, row_v, factx->DV); + } + } hmx_interleave_cols_to_tiles(v_tiles_dst, (const __fp16 *) args->v_src, total_rows, factx->DV, args->src_stride, (uint32_t) args->n_col_tiles, start, end); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, (uint16_t) (args->kv_start + start)); @@ -1809,6 +1845,8 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { memset(&factx, 0, sizeof(factx)); factx.octx = octx; factx.sinks = octx->src[4]; // NULL if this op has no attention sinks + factx.k = k; + factx.v = v; factx.n_threads = kparams->n_threads; factx.DK = DK; factx.DV = DV; @@ -1853,10 +1891,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { // ======== VTCM allocation (GQA-aware) ======== // K/V row sizes drive the DMA descriptors (not the VTCM layout) and are used // throughout the KV loop below. - const size_t size_k_row = DK * sizeof(__fp16); - const size_t size_v_row = DV * sizeof(__fp16); - const size_t size_k_row_padded = hex_round_up(size_k_row, 128); - const size_t size_v_row_padded = hex_round_up(size_v_row, 128); + const size_t size_k_row = htp_tensor_get_row_size(k->type, DK); + const size_t size_v_row = htp_tensor_get_row_size(v->type, DV); + const size_t size_k_row_padded = hex_round_up(DK * sizeof(__fp16), 128); + const size_t size_v_row_padded = hex_round_up(DV * sizeof(__fp16), 128); // Build the VTCM layout once (shared with the host estimator) and place every // scratch buffer at its computed offset. @@ -2348,7 +2386,9 @@ int op_flash_attn_ext(struct htp_ops_context * octx) { const struct htp_tensor * dst = octx->dst; // Check support - if ((q->type != HTP_TYPE_F16 && q->type != HTP_TYPE_F32) || k->type != HTP_TYPE_F16 || v->type != HTP_TYPE_F16) { + if ((q->type != HTP_TYPE_F16 && q->type != HTP_TYPE_F32) || + (k->type != HTP_TYPE_F16 && k->type != HTP_TYPE_Q8_0) || + (v->type != HTP_TYPE_F16 && v->type != HTP_TYPE_Q8_0)) { return HTP_STATUS_NO_SUPPORT; } @@ -2364,6 +2404,8 @@ int op_flash_attn_ext(struct htp_ops_context * octx) { struct htp_fa_context factx; factx.octx = octx; + factx.k = k; + factx.v = v; factx.t_start = HAP_perf_get_qtimer_count(); diff --git a/ggml/src/ggml-hexagon/htp/get-rows-ops.c b/ggml/src/ggml-hexagon/htp/get-rows-ops.c index bf7063e9880..05769d17f74 100644 --- a/ggml/src/ggml-hexagon/htp/get-rows-ops.c +++ b/ggml/src/ggml-hexagon/htp/get-rows-ops.c @@ -12,18 +12,17 @@ #include "ggml-common.h" #include "htp-ctx.h" #include "htp-ops.h" -#include "htp-ops.h" +#include "htp-tensor.h" #include "hvx-utils.h" +#include "hvx-quant.h" +#include "get-rows-ops.h" +#include "work-queue.h" struct get_rows_context { struct htp_ops_context * octx; - uint32_t tasks_per_thread; - uint32_t total_tasks; - uint32_t chunks_per_row; - uint32_t chunk_size; - struct fastdiv_values get_rows_div_ne10; - struct fastdiv_values get_rows_div_ne10_ne11; - struct fastdiv_values get_rows_div_chunks_per_row; + const struct htp_get_rows_kernel_params * kparams; + struct htp_get_rows_vtcm_layout vtcm_layout; + uint8_t * vtcm_base; }; #define get_rows_preamble \ @@ -56,102 +55,161 @@ struct get_rows_context { \ const uint32_t nr = ne10 * ne11 * ne12; -static void get_rows_thread_f32_f32_dma(unsigned int nth, unsigned int ith, void *data) { - struct get_rows_context * grctx = (struct get_rows_context *)data; - struct htp_ops_context * octx = grctx->octx; - get_rows_preamble; - - uint64_t qt = HAP_perf_get_qtimer_count(); - - const uint32_t dr = grctx->tasks_per_thread; - const uint32_t ir0 = dr * ith; - if (ir0 >= grctx->total_tasks) { - return; - } - const uint32_t ir1 = MIN(ir0 + dr, grctx->total_tasks); - - const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32); - - dma_queue * dma_queue = octx->ctx->dma[ith]; - for (uint32_t i = ir0; i < ir1; ++i) { - const uint32_t i12 = fastdiv(i, &grctx->get_rows_div_ne10_ne11); - const uint32_t rem = i - i12 * ne11 * ne10; - const uint32_t i11 = fastdiv(rem, &grctx->get_rows_div_ne10); - const uint32_t i10 = rem - i11 * ne10; - - const uintptr_t src1_addr = octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12; - uint32_t i01 = is_i32 ? *(int32_t *)src1_addr : *(int64_t *)src1_addr; - - if (i01 >= ne01) { - continue; - } - - const uintptr_t src0_ptr = octx->src[0]->data + i01*nb01 + i11*nb02 + i12*nb03; - const uintptr_t dst_ptr = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3; - - while (!dma_queue_push(dma_queue, dma_make_ptr((void *)dst_ptr, (const void *)src0_ptr), nb1, nb01, ne00 * sizeof(float), 1)) { - dma_queue_pop(dma_queue); - } - } - dma_queue_flush(dma_queue); - - qt = HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - qt); - FARF(HIGH, "get-rows-f32-f32-dma %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u usec %u\n", ith, nth, - ne00, ne01, ne02, ne03, ir0, ir1, ne10, ne11, ne12, ne13, ne0, ne1, ne2, ne3, (unsigned) qt); +#define GET_ROWS_THREAD_ST_FN(IDX_TYPE) \ +static void get_rows_thread_st_##IDX_TYPE(unsigned int nth, unsigned int ith, void *data) { \ + struct get_rows_context * grctx = (struct get_rows_context *)data; \ + struct htp_ops_context * octx = grctx->octx; \ + const struct htp_get_rows_kernel_params * kparams = grctx->kparams; \ + get_rows_preamble; \ + const uint32_t dr = kparams->tasks_per_thread; \ + const uint32_t ir0 = dr * ith; \ + if (ir0 >= kparams->total_tasks) { \ + return; \ + } \ + const uint32_t ir1 = MIN(ir0 + dr, kparams->total_tasks); \ + const uint32_t row_size_bytes = htp_tensor_get_row_size(octx->src[0]->type, ne00); \ + dma_queue * dma_queue = octx->ctx->dma[ith]; \ + for (uint32_t i = ir0; i < ir1; ++i) { \ + const uint32_t i12 = fastdiv(i, &kparams->div_ne10_ne11); \ + const uint32_t rem = i - i12 * ne11 * ne10; \ + const uint32_t i11 = fastdiv(rem, &kparams->div_ne10); \ + const uint32_t i10 = rem - i11 * ne10; \ + const IDX_TYPE * src1_ptr = (const IDX_TYPE *)(octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12); \ + const uint32_t i01 = (uint32_t)*src1_ptr; \ + assert(i01 < ne01); \ + const uint32_t q02 = fastdiv(i11, &kparams->div_ne02); \ + const uint32_t i02 = i11 - q02 * ne02; \ + const uint32_t q03 = fastdiv(i12, &kparams->div_ne03); \ + const uint32_t i03 = i12 - q03 * ne03; \ + const uintptr_t src0_ptr = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03; \ + const uintptr_t dst_ptr = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3; \ + while (!dma_queue_push(dma_queue, dma_make_ptr((void *)dst_ptr, (const void *)src0_ptr), nb1, nb01, \ + row_size_bytes, 1)) { \ + dma_queue_pop(dma_queue); \ + } \ + } \ + dma_queue_flush(dma_queue); \ } -static void get_rows_thread_f32_f32_hvx(unsigned int nth, unsigned int ith, void *data) { - struct get_rows_context * grctx = (struct get_rows_context *)data; - struct htp_ops_context * octx = grctx->octx; - get_rows_preamble; - - uint64_t qt = HAP_perf_get_qtimer_count(); - - const uint32_t dr = grctx->tasks_per_thread; - const uint32_t ir0 = dr * ith; - if (ir0 >= grctx->total_tasks) { - return; - } - const uint32_t ir1 = MIN(ir0 + dr, grctx->total_tasks); - - const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32); +GET_ROWS_THREAD_ST_FN(int32_t) +GET_ROWS_THREAD_ST_FN(int64_t) + +#define GET_ROWS_THREAD_DT_FN(TYPE_NAME, SRC0_SIZE_EXPR, IDX_TYPE, COMPUTE_EXPR) \ +static void get_rows_thread_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsigned int ith, void *data) { \ + struct get_rows_context * grctx = (struct get_rows_context *)data; \ + struct htp_ops_context * octx = grctx->octx; \ + const struct htp_get_rows_kernel_params * kparams = grctx->kparams; \ + get_rows_preamble; \ + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ + const uint32_t dr = kparams->tasks_per_thread; \ + const uint32_t ir0 = dr * ith; \ + if (ir0 >= kparams->total_tasks) { \ + return; \ + } \ + const uint32_t ir1 = MIN(ir0 + dr, kparams->total_tasks); \ + const uint32_t chunks_per_row = kparams->chunks_per_row; \ + const uint32_t chunk_size = kparams->chunk_size; \ + dma_queue * dma_queue = octx->ctx->dma[ith]; \ + const struct htp_get_rows_vtcm_layout * vtcm_layout = &grctx->vtcm_layout; \ + uint8_t * vtcm_src0 = grctx->vtcm_base + vtcm_layout->off_src0 + ith * vtcm_layout->src0_bytes_per_thread; \ + uint8_t * vtcm_dst = grctx->vtcm_base + vtcm_layout->off_dst + ith * vtcm_layout->dst_bytes_per_thread; \ + for (uint32_t step = 0, spad_idx = 0; step < ir1 - ir0 && spad_idx < 2; ++step, spad_idx++) { \ + const uint32_t i = ir0 + step; \ + const uint32_t row_idx = fastdiv(i, &kparams->div_chunks_per_row); \ + const uint32_t chunk_idx = i - row_idx * chunks_per_row; \ + const uint32_t i12 = fastdiv(row_idx, &kparams->div_ne10_ne11); \ + const uint32_t rem = row_idx - i12 * ne11 * ne10; \ + const uint32_t i11 = fastdiv(rem, &kparams->div_ne10); \ + const uint32_t i10 = rem - i11 * ne10; \ + const IDX_TYPE * src1_ptr = (const IDX_TYPE *)(octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12); \ + const uint32_t i01 = (uint32_t)*src1_ptr; \ + assert(i01 < ne01); \ + const uint32_t q02 = fastdiv(i11, &kparams->div_ne02); \ + const uint32_t i02 = i11 - q02 * ne02; \ + const uint32_t q03 = fastdiv(i12, &kparams->div_ne03); \ + const uint32_t i03 = i12 - q03 * ne03; \ + const uint32_t offset = chunk_idx * chunk_size; \ + const uint32_t cur_elems = (offset < ne00) ? MIN(chunk_size, ne00 - offset) : 0; \ + const uint32_t cur_src0_bytes = SRC0_SIZE_EXPR(cur_elems); \ + const uint32_t cur_dst_bytes = cur_elems * sizeof(float); \ + const uintptr_t src0_ptr = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03 + SRC0_SIZE_EXPR(offset); \ + dma_queue_push(dma_queue, \ + dma_make_ptr((void *)(uintptr_t)octx->dst->data, \ + vtcm_dst + spad_idx * vtcm_layout->dst_spad_half_size), \ + cur_dst_bytes, vtcm_layout->dst_spad_half_size, cur_dst_bytes, 0); \ + dma_queue_push(dma_queue, \ + dma_make_ptr((void *)(vtcm_src0 + spad_idx * vtcm_layout->src0_spad_half_size), \ + (const void *)src0_ptr), \ + vtcm_layout->src0_spad_half_size, cur_src0_bytes, cur_src0_bytes, 1); \ + } \ + for (uint32_t step = 0; step < ir1 - ir0; ++step) { \ + const uint32_t i = ir0 + step; \ + void * dst_spad = (void *) dma_queue_pop(dma_queue).src; \ + void * src_spad = (void *) dma_queue_pop(dma_queue).dst; \ + const uint32_t row_idx = fastdiv(i, &kparams->div_chunks_per_row); \ + const uint32_t chunk_idx = i - row_idx * chunks_per_row; \ + const uint32_t i12 = fastdiv(row_idx, &kparams->div_ne10_ne11); \ + const uint32_t rem = row_idx - i12 * ne11 * ne10; \ + const uint32_t i11 = fastdiv(rem, &kparams->div_ne10); \ + const uint32_t i10 = rem - i11 * ne10; \ + const uint32_t offset = chunk_idx * chunk_size; \ + const uint32_t cur_elems = (offset < ne00) ? MIN(chunk_size, ne00 - offset) : 0; \ + const uint32_t cur_dst_bytes = cur_elems * sizeof(float); \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, i); \ + COMPUTE_EXPR; \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, i); \ + const uintptr_t dst_ptr = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3 + offset * sizeof(float); \ + dma_queue_push(dma_queue, \ + dma_make_ptr((void *)dst_ptr, (const void *)dst_spad), \ + cur_dst_bytes, vtcm_layout->dst_spad_half_size, cur_dst_bytes, 1); \ + const uint32_t next_step = step + 2; \ + if (next_step < ir1 - ir0) { \ + const uint32_t pi = ir0 + next_step; \ + const uint32_t prow_idx = fastdiv(pi, &kparams->div_chunks_per_row); \ + const uint32_t pchunk_idx = pi - prow_idx * chunks_per_row; \ + const uint32_t pi12 = fastdiv(prow_idx, &kparams->div_ne10_ne11); \ + const uint32_t prem = prow_idx - pi12 * ne11 * ne10; \ + const uint32_t pi11 = fastdiv(prem, &kparams->div_ne10); \ + const uint32_t pi10 = prem - pi11 * ne10; \ + const IDX_TYPE * psrc1_ptr = (const IDX_TYPE *)(octx->src[1]->data + pi10*nb10 + pi11*nb11 + pi12*nb12); \ + const uint32_t pi01 = (uint32_t)*psrc1_ptr; \ + assert(pi01 < ne01); \ + const uint32_t pq02 = fastdiv(pi11, &kparams->div_ne02); \ + const uint32_t pi02 = pi11 - pq02 * ne02; \ + const uint32_t pq03 = fastdiv(pi12, &kparams->div_ne03); \ + const uint32_t pi03 = pi12 - pq03 * ne03; \ + const uint32_t poffset = pchunk_idx * chunk_size; \ + const uint32_t pcur_elems = (poffset < ne00) ? MIN(chunk_size, ne00 - poffset) : 0; \ + const uint32_t pcur_src0_bytes = SRC0_SIZE_EXPR(pcur_elems); \ + const uintptr_t psrc0_ptr = \ + octx->src[0]->data + pi01*nb01 + pi02*nb02 + pi03*nb03 + SRC0_SIZE_EXPR(poffset); \ + dma_queue_push(dma_queue, \ + dma_make_ptr((void *)src_spad, (const void *)psrc0_ptr), \ + vtcm_layout->src0_spad_half_size, pcur_src0_bytes, pcur_src0_bytes, 1); \ + } \ + } \ + dma_queue_flush(dma_queue); \ +} - const uint32_t chunks_per_row = grctx->chunks_per_row; - const uint32_t chunk_size = grctx->chunk_size; - for (uint32_t i = ir0; i < ir1; ++i) { - const uint32_t row_idx = fastdiv(i, &grctx->get_rows_div_chunks_per_row); - const uint32_t chunk_idx = i - row_idx * chunks_per_row; +#define F32_BYTES(n) ((n) * sizeof(float)) +#define F16_BYTES(n) ((n) * sizeof(__fp16)) +#define Q8_0_BYTES(n) (((n) / 32) * sizeof(block_q8_0)) - const uint32_t i12 = fastdiv(row_idx, &grctx->get_rows_div_ne10_ne11); - const uint32_t rem = row_idx - i12 * ne11 * ne10; - const uint32_t i11 = fastdiv(rem, &grctx->get_rows_div_ne10); - const uint32_t i10 = rem - i11 * ne10; +GET_ROWS_THREAD_DT_FN(f32, F32_BYTES, int32_t, { if (cur_elems > 0) hvx_copy_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, cur_elems); }) +GET_ROWS_THREAD_DT_FN(f32, F32_BYTES, int64_t, { if (cur_elems > 0) hvx_copy_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, cur_elems); }) - const uintptr_t src1_addr = octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12; - uint32_t i01 = is_i32 ? *(int32_t *)src1_addr : *(int64_t *)src1_addr; +GET_ROWS_THREAD_DT_FN(f16, F16_BYTES, int32_t, { hvx_dequantize_row_f16_f32((float *)dst_spad, src_spad, ne00); }) +GET_ROWS_THREAD_DT_FN(f16, F16_BYTES, int64_t, { hvx_dequantize_row_f16_f32((float *)dst_spad, src_spad, ne00); }) - if (i01 >= ne01) { - continue; - } - - const uint32_t offset = chunk_idx * chunk_size; - if (offset < ne00) { - const uint32_t copy_size = MIN(chunk_size, ne00 - offset); - const uintptr_t src0_ptr = octx->src[0]->data + i01*nb01 + i11*nb02 + i12*nb03 + offset * sizeof(float); - const uintptr_t dst_ptr = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3 + offset * sizeof(float); - hvx_copy_f32_uu((uint8_t *)dst_ptr, (const uint8_t *)src0_ptr, copy_size); - } - } - - qt = HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - qt); - FARF(HIGH, "get-rows-f32-f32-hvx %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u usec %u\n", ith, nth, - ne00, ne01, ne02, ne03, ir0, ir1, ne10, ne11, ne12, ne13, ne0, ne1, ne2, ne3, (unsigned) qt); -} +GET_ROWS_THREAD_DT_FN(q8_0, Q8_0_BYTES, int32_t, { hvx_dequantize_row_q8_0_f32((float *)dst_spad, src_spad, ne00); }) +GET_ROWS_THREAD_DT_FN(q8_0, Q8_0_BYTES, int64_t, { hvx_dequantize_row_q8_0_f32((float *)dst_spad, src_spad, ne00); }) int op_get_rows(struct htp_ops_context * octx) { - get_rows_preamble; + const struct htp_get_rows_kernel_params * kparams = (const struct htp_get_rows_kernel_params *) octx->kernel_params; - if (octx->src[0]->type != HTP_TYPE_F32) { + if (octx->src[0]->type != HTP_TYPE_F32 && + octx->src[0]->type != HTP_TYPE_F16 && + octx->src[0]->type != HTP_TYPE_Q8_0) { return HTP_STATUS_NO_SUPPORT; } @@ -167,52 +225,28 @@ int op_get_rows(struct htp_ops_context * octx) { return HTP_STATUS_OK; } - const uint32_t nb00 = octx->src[0]->nb[0]; - const uint32_t nb0 = octx->dst->nb[0]; - - const bool can_use_dma = (nb00 == sizeof(float)) && (nb0 == sizeof(float)); - const bool use_dma = can_use_dma && (ne00 >= 2048); - struct get_rows_context grctx; grctx.octx = octx; - grctx.get_rows_div_ne10 = init_fastdiv_values(octx->src[1]->ne[0]); - grctx.get_rows_div_ne10_ne11 = init_fastdiv_values(octx->src[1]->ne[0] * octx->src[1]->ne[1]); + grctx.kparams = kparams; + grctx.vtcm_base = (uint8_t *)octx->ctx->vtcm_base; - if (use_dma) { - grctx.chunks_per_row = 1; - grctx.chunk_size = ne00; - grctx.total_tasks = nr; - grctx.get_rows_div_chunks_per_row = init_fastdiv_values(1); + const uint32_t ne00 = octx->src[0]->ne[0]; + htp_get_rows_vtcm_layout_build(&grctx.vtcm_layout, octx->src[0]->type, ne00, kparams->n_threads); - const uint32_t n_threads = MIN(nr, octx->n_threads); - grctx.tasks_per_thread = (nr + n_threads - 1) / n_threads; + const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32); - worker_pool_run_func(octx->ctx->worker_pool, get_rows_thread_f32_f32_dma, &grctx, n_threads); + work_queue_func_t q_func = NULL; + if (kparams->use_dma) { + q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_st_int32_t : get_rows_thread_st_int64_t); } else { - uint32_t chunks_per_row = 1; - uint32_t chunk_size = ne00; - uint32_t total_tasks = nr; - - if (nr < octx->n_threads) { - const uint32_t min_chunk_size = 1024; - uint32_t max_chunks = ne00 / min_chunk_size; - if (max_chunks == 0) { - max_chunks = 1; - } - chunks_per_row = MIN((octx->n_threads + nr - 1) / nr, max_chunks); - chunk_size = (ne00 + chunks_per_row - 1) / chunks_per_row; - total_tasks = nr * chunks_per_row; + switch (octx->src[0]->type) { + case HTP_TYPE_F32: q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_f32_int32_t : get_rows_thread_f32_int64_t); break; + case HTP_TYPE_F16: q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_f16_int32_t : get_rows_thread_f16_int64_t); break; + case HTP_TYPE_Q8_0: q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_q8_0_int32_t : get_rows_thread_q8_0_int64_t); break; + default: return HTP_STATUS_NO_SUPPORT; } - - grctx.chunks_per_row = chunks_per_row; - grctx.chunk_size = chunk_size; - grctx.total_tasks = total_tasks; - grctx.get_rows_div_chunks_per_row = init_fastdiv_values(chunks_per_row); - - const uint32_t n_threads = MIN(total_tasks, octx->n_threads); - grctx.tasks_per_thread = (total_tasks + n_threads - 1) / n_threads; - - worker_pool_run_func(octx->ctx->worker_pool, get_rows_thread_f32_f32_hvx, &grctx, n_threads); } + + work_queue_run(octx->ctx->work_queue, q_func, &grctx, kparams->n_threads); return HTP_STATUS_OK; } diff --git a/ggml/src/ggml-hexagon/htp/get-rows-ops.h b/ggml/src/ggml-hexagon/htp/get-rows-ops.h new file mode 100644 index 00000000000..0e7c2ca8cf0 --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/get-rows-ops.h @@ -0,0 +1,77 @@ +#ifndef HTP_GET_ROWS_OPS_H +#define HTP_GET_ROWS_OPS_H + +#include "hex-fastdiv.h" + +struct htp_get_rows_kernel_params { + int32_t n_threads; + int32_t use_dma; + int32_t chunks_per_row; + int32_t chunk_size; + int32_t total_tasks; + int32_t tasks_per_thread; + int32_t vtcm_size; + + // Fastdiv helpers + struct fastdiv_values div_ne10; + struct fastdiv_values div_ne10_ne11; + struct fastdiv_values div_chunks_per_row; + struct fastdiv_values div_ne02; + struct fastdiv_values div_ne03; +}; + +struct htp_get_rows_vtcm_layout { + size_t total_bytes; + size_t off_src0; + size_t off_dst; + + size_t src0_bytes_per_thread; + size_t dst_bytes_per_thread; + + size_t src0_spad_half_size; + size_t dst_spad_half_size; +}; + +static inline void htp_get_rows_vtcm_layout_build( + struct htp_get_rows_vtcm_layout * vtcm_layout, + int type, + uint32_t ne00, + uint32_t n_threads) { + + uint32_t src0_row_size = 0; + switch (type) { + case 0: // HTP_TYPE_F32 + src0_row_size = ne00 * 4; + break; + case 1: // HTP_TYPE_F16 + src0_row_size = ne00 * 2; + break; + case 8: // HTP_TYPE_Q8_0 + src0_row_size = (ne00 / 32) * 34; + break; + default: + src0_row_size = 0; + break; + } + + size_t src0_row_size_aligned = (src0_row_size + 255) & ~255; + size_t dst_row_size_aligned = (ne00 * sizeof(float) + 255) & ~255; + + vtcm_layout->src0_spad_half_size = src0_row_size_aligned; + vtcm_layout->dst_spad_half_size = dst_row_size_aligned; + + vtcm_layout->src0_bytes_per_thread = src0_row_size_aligned * 2; + vtcm_layout->dst_bytes_per_thread = dst_row_size_aligned * 2; + + vtcm_layout->off_src0 = 0; + vtcm_layout->off_dst = vtcm_layout->off_src0 + vtcm_layout->src0_bytes_per_thread * n_threads; + vtcm_layout->total_bytes = vtcm_layout->off_dst + vtcm_layout->dst_bytes_per_thread * n_threads; +} + +#if defined(__cplusplus) +static_assert(sizeof(struct htp_get_rows_kernel_params) <= 128, "htp_get_rows_kernel_params is too large for kernel_params blob"); +#else +_Static_assert(sizeof(struct htp_get_rows_kernel_params) <= 128, "htp_get_rows_kernel_params is too large for kernel_params blob"); +#endif + +#endif // HTP_GET_ROWS_OPS_H diff --git a/ggml/src/ggml-hexagon/htp/hex-utils.h b/ggml/src/ggml-hexagon/htp/hex-utils.h index 93e87efcb4c..1b396503000 100644 --- a/ggml/src/ggml-hexagon/htp/hex-utils.h +++ b/ggml/src/ggml-hexagon/htp/hex-utils.h @@ -39,17 +39,22 @@ static inline void hex_l2fetch_block(const void * addr, size_t size) { #define HEX_L2_LINE_SIZE 128 #define HEX_L2_BLOCK_SIZE (HEX_L2_LINE_SIZE * 4) // flush granularity (lines per loop iteration) +#define HEX_L2_FLUSH_IL_THRESHOLD 1024 // inline flush threshold #define HEX_L2_FLUSH_WQ_THRESHOLD (4 * 1024) #define HEX_L2_FLUSH_ALL_THRESHOLD (4 * 1024 * 1024) static inline void hex_l2flush(void * addr, size_t size) { const uint32_t s = ((uint32_t) addr) & ~(HEX_L2_LINE_SIZE - 1); const uint32_t e = (((uint32_t) addr) + size + HEX_L2_LINE_SIZE - 1) & ~(HEX_L2_LINE_SIZE - 1); - for (uint32_t i = s; i < e; i += HEX_L2_BLOCK_SIZE) { - Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 0); - Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 1); - Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 2); - Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 3); + const uint32_t eb = s + ((e - s) & ~(HEX_L2_BLOCK_SIZE - 1)); + for (uint32_t i = s; i < eb; i += HEX_L2_BLOCK_SIZE) { + Q6_dccleaninva_A((void *) (i + HEX_L2_LINE_SIZE * 0)); + Q6_dccleaninva_A((void *) (i + HEX_L2_LINE_SIZE * 1)); + Q6_dccleaninva_A((void *) (i + HEX_L2_LINE_SIZE * 2)); + Q6_dccleaninva_A((void *) (i + HEX_L2_LINE_SIZE * 3)); + } + for (uint32_t i = eb; i < e; i += HEX_L2_LINE_SIZE) { + Q6_dccleaninva_A((void *) i); } } diff --git a/ggml/src/ggml-hexagon/htp/htp-ctx.h b/ggml/src/ggml-hexagon/htp/htp-ctx.h index e0f9a0c40d1..88ecf144b94 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ctx.h +++ b/ggml/src/ggml-hexagon/htp/htp-ctx.h @@ -117,8 +117,7 @@ struct htp_context { int op_matmul(struct htp_ops_context * octx); int op_matmul_id(struct htp_ops_context * octx); -int op_matmul_qkv(struct htp_ops_context * octx); -int op_matmul_ffn(struct htp_ops_context * octx); +int op_matmul_nx(struct htp_ops_context * octx); int op_binary(struct htp_ops_context * octx); int op_unary(struct htp_ops_context * octx); int op_sum_rows(struct htp_ops_context * octx); @@ -141,5 +140,6 @@ int op_solve_tri(struct htp_ops_context * octx); int op_gated_delta_net(struct htp_ops_context * octx); int op_pad(struct htp_ops_context * octx); int op_im2col(struct htp_ops_context * octx); +int op_allreduce(struct htp_ops_context * octx); #endif /* HTP_CTX_H */ diff --git a/ggml/src/ggml-hexagon/htp/htp-ops.h b/ggml/src/ggml-hexagon/htp/htp-ops.h index a138f062aa6..e804844d599 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ops.h +++ b/ggml/src/ggml-hexagon/htp/htp-ops.h @@ -43,13 +43,6 @@ enum htp_data_type { -// Mask to enable various stages of the Ops. -// Used for debugging and profiling. -enum htp_op_stage { - HTP_OPSTAGE_QUEUE = (1 << 0), // Enable Queueing (ie calls into NPU) - HTP_OPSTAGE_COMPUTE = (1 << 1), // Enable Compute -}; - // Do not reorder first 4 (used as an index) enum htp_op_code { HTP_OP_MUL = 0, @@ -58,8 +51,7 @@ enum htp_op_code { HTP_OP_DIV = 3, HTP_OP_MUL_MAT, HTP_OP_MUL_MAT_ID, - HTP_OP_MUL_MAT_QKV, - HTP_OP_MUL_MAT_FFN, + HTP_OP_MUL_MAT_NX, HTP_OP_MUL_MAT_ADD, HTP_OP_RMS_NORM, HTP_OP_RMS_NORM_MUL, @@ -70,6 +62,8 @@ enum htp_op_code { HTP_OP_UNARY_NEG, HTP_OP_UNARY_SOFTPLUS, HTP_OP_UNARY_TANH, + HTP_OP_UNARY_ABS, + HTP_OP_UNARY_LOG, HTP_OP_GLU_SWIGLU, HTP_OP_GLU_SWIGLU_OAI, HTP_OP_GLU_GEGLU, @@ -99,12 +93,15 @@ enum htp_op_code { HTP_OP_CONCAT, HTP_OP_CLAMP, HTP_OP_IM2COL, + HTP_OP_FENCE, + HTP_OP_ALLREDUCE, + HTP_OP_ALLREDUCE_ADD, HTP_OP_INVALID }; #define HTP_OP_MAX_DIMS 4 // aka GGML_MAX_DIMS -#define HTP_OP_MAX_INPUTS 6 // aka GGML_MAX_SRCS +#define HTP_OP_MAX_INPUTS 10 // aka GGML_MAX_SRCS #define HTP_OP_MAX_OUTPUTS 4 #define HTP_OP_MAX_PARAMS 16 // aka GGML_MAX_OP_PARAMS #define HTP_OP_MAX_KERN_PARAMS 32 @@ -112,13 +109,16 @@ enum htp_op_code { #define HTP_OP_MAX_BUFS 16 #define HTP_OP_MAX_TENSORS 8192 // must stay under 64K (uint16) +#define HTP_FENCE_TIMEOUT (1000000000ULL) + #define HTP_OP_MAX_VMEM_DEFAULT (3355443200u) #define HTP_MMAP_MAX_VMEM (2147483648u) enum htp_tensor_flags { - HTP_TENSOR_COMPUTE = (1U << 0), // Tensor buffer temporal compute data (not weights) - HTP_TENSOR_DIRTY = (1U << 1) // Tensor buffer is dirty and needs to be flushed + HTP_TENSOR_WEIGHT = (1U << 0), // Tensor buffer model weight data (not compute) + HTP_TENSOR_REPACK = (1U << 1), // Tensor is in repacked tiled format + HTP_TENSOR_FENCE = (1U << 2) // Tensor is synchronization fence (explicitly managed) }; // Tensor descriptor @@ -175,6 +175,7 @@ enum htp_trace_event_id { HTP_TRACE_EVT_L2FLUSH = 1, HTP_TRACE_EVT_INIT = 2, HTP_TRACE_EVT_BUFF = 3, + HTP_TRACE_EVT_FENCE = 4, HTP_TRACE_EVT_HVX_COMP = 20, HTP_TRACE_EVT_HVX_A_QUANT = 21, @@ -215,6 +216,7 @@ struct htp_opbatch_req { uint32_t n_ops; // Number of ops uint32_t n_traces; // Number of trace descriptors per thread uint32_t pad; // unused + uint64_t seq; // Sequence number // struct htp_buf_desc bufs[]; -- dspqueue buf 0 // struct htp_tensor tensors[]; -- dspqueue buf 0 // struct htp_op_desc ops[]; -- dspqueue buf 0 @@ -231,6 +233,7 @@ struct htp_opbatch_rsp { uint32_t pad; // align to 8 bytes uint64_t cycles_start; // Start cycle counter uint64_t cycles_stop; // Stop cycle counter + uint64_t seq; // Sequence number // struct htp_prof_desc profs[]; -- dspqueue buf 0 }; diff --git a/ggml/src/ggml-hexagon/htp/htp-tensor.c b/ggml/src/ggml-hexagon/htp/htp-tensor.c index 39436e26dff..ae377c9221f 100644 --- a/ggml/src/ggml-hexagon/htp/htp-tensor.c +++ b/ggml/src/ggml-hexagon/htp/htp-tensor.c @@ -79,7 +79,14 @@ void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * co for (uint32_t i = 0; i < n; i++) { const struct htp_tensor * t = tensors[i]; - if (!t) continue; + if (!t || (t->flags & (HTP_TENSOR_WEIGHT | HTP_TENSOR_FENCE))) { + continue; + } + + if (t->size <= HEX_L2_FLUSH_IL_THRESHOLD) { + hex_l2flush((void *) (uintptr_t) t->data, t->size); + continue; + } uint32_t t_start = t->data; uint32_t t_end = t_start + t->size; @@ -242,7 +249,7 @@ void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * co for (uint32_t i = 0; i < n; i++) { const struct htp_tensor * t = tensors[i]; - if (t && (t->flags & HTP_TENSOR_COMPUTE) && is_tensor_dirty(ctx, t)) { + if (t && !(t->flags & (HTP_TENSOR_WEIGHT | HTP_TENSOR_FENCE)) && is_tensor_dirty(ctx, t)) { dirty_tensors[n_dirty++] = t; total_dirty += t->size; } diff --git a/ggml/src/ggml-hexagon/htp/htp-tensor.h b/ggml/src/ggml-hexagon/htp/htp-tensor.h index 2c3fc54c748..c9cadbae3f2 100644 --- a/ggml/src/ggml-hexagon/htp/htp-tensor.h +++ b/ggml/src/ggml-hexagon/htp/htp-tensor.h @@ -13,6 +13,15 @@ static inline uint32_t * htp_tensor_flags(const struct htp_tensor * t) { return (uint32_t *) &t->flags; } +static inline uint32_t htp_tensor_get_row_size(int type, uint32_t ne00) { + switch (type) { + case HTP_TYPE_F32: return ne00 * 4; + case HTP_TYPE_F16: return ne00 * 2; + case HTP_TYPE_Q8_0: return (ne00 / 32) * 34; + default: return 0; + } +} + struct htp_context; void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n); void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n); diff --git a/ggml/src/ggml-hexagon/htp/hvx-arith.h b/ggml/src/ggml-hexagon/htp/hvx-arith.h index 82e3416970b..765c3577668 100644 --- a/ggml/src/ggml-hexagon/htp/hvx-arith.h +++ b/ggml/src/ggml-hexagon/htp/hvx-arith.h @@ -17,9 +17,9 @@ #define hvx_arith_loop_body(dst_type, src0_type, src1_type, elem_size, vec_store, vec_op) \ do { \ - dst_type * restrict vdst = (dst_type *) dst; \ - src0_type * restrict vsrc0 = (src0_type *) src0; \ - src1_type * restrict vsrc1 = (src1_type *) src1; \ + dst_type * vdst = (dst_type *) dst; \ + src0_type * vsrc0 = (src0_type *) src0; \ + src1_type * vsrc1 = (src1_type *) src1; \ \ const uint32_t epv = 128 / (elem_size); \ const uint32_t nvec = n / epv; \ @@ -57,40 +57,40 @@ // Generic macro to define alignment permutations for an op #define DEFINE_HVX_BINARY_OP_VARIANTS(OP_NAME, OP_MACRO, ELEM_TYPE) \ -static inline void OP_NAME##_aaa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ +static inline void OP_NAME##_aaa(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \ assert((uintptr_t) dst % 128 == 0); \ assert((uintptr_t) src0 % 128 == 0); \ assert((uintptr_t) src1 % 128 == 0); \ hvx_arith_loop_body(HVX_Vector, HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \ } \ -static inline void OP_NAME##_aau(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ +static inline void OP_NAME##_aau(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \ assert((uintptr_t) dst % 128 == 0); \ assert((uintptr_t) src0 % 128 == 0); \ hvx_arith_loop_body(HVX_Vector, HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \ } \ -static inline void OP_NAME##_aua(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ +static inline void OP_NAME##_aua(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \ assert((uintptr_t) dst % 128 == 0); \ assert((uintptr_t) src1 % 128 == 0); \ hvx_arith_loop_body(HVX_Vector, HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \ } \ -static inline void OP_NAME##_auu(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ +static inline void OP_NAME##_auu(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \ assert((uintptr_t) dst % 128 == 0); \ hvx_arith_loop_body(HVX_Vector, HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \ } \ -static inline void OP_NAME##_uaa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ +static inline void OP_NAME##_uaa(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \ assert((uintptr_t) src0 % 128 == 0); \ assert((uintptr_t) src1 % 128 == 0); \ hvx_arith_loop_body(HVX_UVector, HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \ } \ -static inline void OP_NAME##_uau(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ +static inline void OP_NAME##_uau(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \ assert((uintptr_t) src0 % 128 == 0); \ hvx_arith_loop_body(HVX_UVector, HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \ } \ -static inline void OP_NAME##_uua(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ +static inline void OP_NAME##_uua(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \ assert((uintptr_t) src1 % 128 == 0); \ hvx_arith_loop_body(HVX_UVector, HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \ } \ -static inline void OP_NAME##_uuu(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ +static inline void OP_NAME##_uuu(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \ hvx_arith_loop_body(HVX_UVector, HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \ } \ @@ -358,6 +358,34 @@ static inline void hvx_clamp_scalar_f32(uint8_t * restrict dst, const uint8_t * } } +// +// Abs +// + +static inline void hvx_abs_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) dst % 128 == 0); + assert((unsigned long) src % 128 == 0); + + HVX_Vector * restrict vdst = (HVX_Vector *) dst; + HVX_Vector * restrict vsrc = (HVX_Vector *) src; + + const uint32_t elem_size = sizeof(float); + const uint32_t epv = 128 / elem_size; + const uint32_t nvec = n / epv; + const uint32_t nloe = n % epv; + + uint32_t i = 0; + + _Pragma("unroll(4)") + for (; i < nvec; i++) { + vdst[i] = hvx_vec_abs_f32(vsrc[i]); + } + if (nloe) { + HVX_Vector v = hvx_vec_abs_f32(vsrc[i]); + hvx_vec_store_a((void *) &vdst[i], nloe * elem_size, v); + } +} + // // Square // diff --git a/ggml/src/ggml-hexagon/htp/hvx-log.h b/ggml/src/ggml-hexagon/htp/hvx-log.h index 7013dae785a..a209f88d555 100644 --- a/ggml/src/ggml-hexagon/htp/hvx-log.h +++ b/ggml/src/ggml-hexagon/htp/hvx-log.h @@ -62,4 +62,28 @@ static inline HVX_Vector hvx_vec_log_f32(HVX_Vector x) { return hvx_vec_add_f32_f32(term_e, res); } +static inline void hvx_log_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) dst % 128 == 0); + assert((unsigned long) src % 128 == 0); + + HVX_Vector * restrict vdst = (HVX_Vector *) dst; + HVX_Vector * restrict vsrc = (HVX_Vector *) src; + + const uint32_t elem_size = sizeof(float); + const uint32_t epv = 128 / elem_size; + const uint32_t nvec = n / epv; + const uint32_t nloe = n % epv; + + uint32_t i = 0; + + _Pragma("unroll(4)") + for (; i < nvec; i++) { + vdst[i] = hvx_vec_log_f32(vsrc[i]); + } + if (nloe) { + HVX_Vector v = hvx_vec_log_f32(vsrc[i]); + hvx_vec_store_a((void *) &vdst[i], nloe * elem_size, v); + } +} + #endif /* HVX_LOG_H */ diff --git a/ggml/src/ggml-hexagon/htp/hvx-quant.h b/ggml/src/ggml-hexagon/htp/hvx-quant.h new file mode 100644 index 00000000000..6b172cd63c3 --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/hvx-quant.h @@ -0,0 +1,165 @@ +#ifndef HVX_QUANT_H +#define HVX_QUANT_H + +#include +#include +#include + +#include "hvx-arith.h" +#include "hvx-base.h" +#include "hvx-reduce.h" +#include "hvx-repl.h" +#include "hvx-utils.h" + +#ifndef GGML_COMMON_DECL_C +#define GGML_COMMON_DECL_C +#endif +#include "ggml-common.h" +#include "ggml-impl.h" + +static inline void hvx_quantize_row_q8_0_f32(void * restrict dst_ptr, const float * restrict src_ptr, int n) { + const int nb = n / QK8_0; + block_q8_0 * dst = (block_q8_0 *) dst_ptr; + HVX_Vector zero = Q6_V_vzero(); + + int i = 0; + for (; i + 3 < nb; i += 4) { + HVX_Vector * vx = (HVX_Vector *) (src_ptr + i * QK8_0); + + HVX_Vector vmax0_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[0])); + HVX_Vector vmax1_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[1])); + HVX_Vector vmax2_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[2])); + HVX_Vector vmax3_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[3])); + + HVX_Vector vx0_qf = Q6_Vqf32_vsub_VsfVsf(vx[0], zero); + HVX_Vector vx1_qf = Q6_Vqf32_vsub_VsfVsf(vx[1], zero); + HVX_Vector vx2_qf = Q6_Vqf32_vsub_VsfVsf(vx[2], zero); + HVX_Vector vx3_qf = Q6_Vqf32_vsub_VsfVsf(vx[3], zero); + + HVX_Vector vmax0_qf = Q6_Vqf32_vsub_VsfVsf(vmax0_sf, zero); + HVX_Vector vmax1_qf = Q6_Vqf32_vsub_VsfVsf(vmax1_sf, zero); + HVX_Vector vmax2_qf = Q6_Vqf32_vsub_VsfVsf(vmax2_sf, zero); + HVX_Vector vmax3_qf = Q6_Vqf32_vsub_VsfVsf(vmax3_sf, zero); + + HVX_Vector vmax01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax1_qf, vmax0_qf))); + HVX_Vector vmax23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax3_qf, vmax2_qf))); + + HVX_Vector vx01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx1_qf, vx0_qf))); + HVX_Vector vx23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx3_qf, vx2_qf))); + + HVX_Vector vd01_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax01_hf, Q6_Vh_vsplat_R(0x2008)); // 1.0 / 127.0 + HVX_Vector vd23_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax23_hf, Q6_Vh_vsplat_R(0x2008)); // 1.0 / 127.0 + HVX_Vector vd01_hf = Q6_Vhf_equals_Vqf16(vd01_qf16); + HVX_Vector vd23_hf = Q6_Vhf_equals_Vqf16(vd23_qf16); + + HVX_Vector vd01_inv_hf = hvx_vec_inverse_f16(vd01_hf); + HVX_Vector vd23_inv_hf = hvx_vec_inverse_f16(vd23_hf); + vx01_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx01_hf, vd01_inv_hf)); + vx23_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx23_hf, vd23_inv_hf)); + + HVX_Vector vx01_i16 = hvx_vec_i16_from_hf_rnd_sat(vx01_hf); + HVX_Vector vx23_i16 = hvx_vec_i16_from_hf_rnd_sat(vx23_hf); + HVX_Vector vx_i8 = Q6_Vb_vpack_VhVh_sat(vx23_i16, vx01_i16); + + hvx_vec_store_u(&dst[i + 0].d, 2, vd01_hf); + hvx_vec_store_u(dst[i + 0].qs, 32, vx_i8); + + hvx_vec_store_u(&dst[i + 1].d, 2, Q6_V_vror_VR(vd01_hf, 64)); + hvx_vec_store_u(dst[i + 1].qs, 32, Q6_V_vror_VR(vx_i8, 32)); + + hvx_vec_store_u(&dst[i + 2].d, 2, vd23_hf); + hvx_vec_store_u(dst[i + 2].qs, 32, Q6_V_vror_VR(vx_i8, 64)); + + hvx_vec_store_u(&dst[i + 3].d, 2, Q6_V_vror_VR(vd23_hf, 64)); + hvx_vec_store_u(dst[i + 3].qs, 32, Q6_V_vror_VR(vx_i8, 96)); + } + + for (; i < nb; i++) { + const float * block_src = src_ptr + i * QK8_0; + HVX_Vector vx = *(const HVX_UVector *) block_src; + HVX_Vector v_abs = hvx_vec_abs_f32(vx); + HVX_Vector v_max = hvx_vec_reduce_max_f32(v_abs); + float amax = hvx_vec_get_f32(v_max); + + const float d = amax / 127.0f; + const float id = d ? (1.0f / d) : 0.0f; + dst[i].d = GGML_FP32_TO_FP16(d); + + HVX_Vector vid = hvx_vec_splat_f32(id); + HVX_Vector v_scaled = hvx_vec_mul_f32_f32(vx, vid); + HVX_Vector v_scaled_qf = Q6_Vqf32_vsub_VsfVsf(v_scaled, zero); + HVX_Vector v_scaled_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(zero, v_scaled_qf))); + HVX_Vector v_i16 = hvx_vec_i16_from_hf_rnd_sat(v_scaled_hf); + HVX_Vector v_i8 = Q6_Vb_vpack_VhVh_sat(zero, v_i16); + + hvx_vec_store_u(dst[i].qs, 32, v_i8); + } +} + +static inline void hvx_dequantize_row_q8_0_f32(float * restrict dst_ptr, const void * restrict src_ptr, int n) { + const int nb = n / QK8_0; + const block_q8_0 * src = (const block_q8_0 *) src_ptr; + + for (int i = 0; i < nb; i++) { + HVX_Vector vd_f16 = Q6_Vh_vsplat_R(*(const int16_t *) &src[i].d); + HVX_VectorPair vp_f32 = hvx_vec_f16_to_f32(vd_f16); + HVX_Vector vd = Q6_V_lo_W(vp_f32); + + HVX_Vector vq_i8 = *(const HVX_UVector *) src[i].qs; + + HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(vq_i8); + HVX_Vector v_i16 = Q6_V_lo_W(p16); + HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(v_i16); + HVX_Vector v_i32 = Q6_V_lo_W(p32); + + HVX_Vector v_f32 = Q6_Vsf_equals_Vw(v_i32); + HVX_Vector res = hvx_vec_mul_f32_f32(v_f32, vd); + + float * block_dst = dst_ptr + i * QK8_0; + hvx_vmem(block_dst) = res; + } +} + +static inline void hvx_dequantize_row_q8_0_f16(__fp16 * restrict dst_ptr, const void * restrict src_ptr, int n) { + const int nb = n / QK8_0; + const block_q8_0 * src = (const block_q8_0 *) src_ptr; + + for (int i = nb - 1; i >= 0; i--) { + HVX_Vector vd_f16 = Q6_Vh_vsplat_R(*(const int16_t *) &src[i].d); + HVX_VectorPair vp_f32 = hvx_vec_f16_to_f32(vd_f16); + HVX_Vector vd = Q6_V_lo_W(vp_f32); + + HVX_Vector vq_i8 = *(const HVX_UVector *) src[i].qs; + + HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(vq_i8); + HVX_Vector v_i16 = Q6_V_lo_W(p16); + HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(v_i16); + HVX_Vector v_i32 = Q6_V_lo_W(p32); + + HVX_Vector v_f32 = Q6_Vsf_equals_Vw(v_i32); + HVX_Vector res_f32 = hvx_vec_mul_f32_f32(v_f32, vd); + + HVX_Vector res_f16 = hvx_vec_f32_to_f16(res_f32, Q6_V_vzero()); + + __fp16 * block_dst = dst_ptr + i * QK8_0; + hvx_vec_store_u(block_dst, QK8_0 * sizeof(__fp16), res_f16); + } +} + +static inline void hvx_dequantize_row_f16_f32(float * restrict dst_ptr, const void * restrict src_ptr, int n) { + const int nb = n / 32; + const _Float16 * src = (const _Float16 *) src_ptr; + + for (int i = 0; i < nb; i++) { + HVX_Vector v_f16 = *(const HVX_UVector *) (src + i * 32); + HVX_VectorPair vp_f32 = hvx_vec_f16_to_f32(v_f16); + HVX_Vector res = Q6_V_lo_W(vp_f32); + + float * block_dst = dst_ptr + i * 32; + hvx_vmem(block_dst) = res; + } +} + + + +#endif // HVX_QUANT_H diff --git a/ggml/src/ggml-hexagon/htp/main.c b/ggml/src/ggml-hexagon/htp/main.c index 880e20c9959..fe7d093a81c 100644 --- a/ggml/src/ggml-hexagon/htp/main.c +++ b/ggml/src/ggml-hexagon/htp/main.c @@ -18,6 +18,7 @@ #include #include #include +#include #include "hex-utils.h" #include "hex-dma.h" @@ -32,6 +33,7 @@ #include "htp_iface.h" #include "work-queue.h" #include "hex-profile.h" +#include "allreduce-ops.h" #define HMX_QUEUE_CAPACITY 16 #define HMX_QUEUE_STACK_SIZE 16384 @@ -46,6 +48,36 @@ struct htp_handle { struct htp_context * ctx; }; +static inline void * htp_mmap(uint32_t fd, uint32_t size) { + void * va = (void *)-1; + for (int retry = 0; retry < 2; retry++) { +#if __HVX_ARCH__ > 73 + va = HAP_mmap2(NULL, size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0); +#else + if (size > HTP_MMAP_MAX_VMEM) { + FARF(ERROR, "mmap failed : size %u exceeds 2GB limit for HAP_mmap", (uint32_t) size); + abort(); + } + va = HAP_mmap(NULL, size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0); +#endif + if (va != (void *)-1 && va != NULL) { + return va; + } + if (retry == 0) { + FARF(HIGH, "mmap failed first try (va %p fd %u size %u), retrying...", va, fd, size); + } + } + return NULL; +} + +static inline void htp_munmap(void * va, uint32_t size) { +#if __HVX_ARCH__ > 73 + HAP_munmap2(va, size); +#else + HAP_munmap(va, size); +#endif +} + AEEResult htp_iface_open(const char * uri, remote_handle64 * handle) { (void) uri; struct htp_handle * h = calloc(1, sizeof(*h)); @@ -127,11 +159,7 @@ AEEResult htp_iface_close(remote_handle64 handle) { // release the mmaps (if any) for (uint32_t i=0; immap[i].size) { -#if __HVX_ARCH__ > 73 - HAP_munmap2((void *) ctx->mmap[i].base, ctx->mmap[i].size); -#else - HAP_munmap((void *) ctx->mmap[i].base, ctx->mmap[i].size); -#endif + htp_munmap((void *) ctx->mmap[i].base, ctx->mmap[i].size); ctx->mmap[i].size = 0; ctx->mmap[i].base = NULL; ctx->mmap[i].fd = -1; @@ -175,18 +203,9 @@ AEEResult htp_iface_mmap(remote_handle64 handle, uint32_t fd, uint32_t size) { struct htp_mmap *m = &ctx->mmap[i]; if (!m->size) { FARF(HIGH, "mmap : fd %u size %u", fd, size); -#if __HVX_ARCH__ > 73 - void *va = HAP_mmap2(NULL, size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0); -#else - if (size > HTP_MMAP_MAX_VMEM) { // HAP_mmap has a size limit of 2GB - FARF(ERROR, "mmap failed : size %u exceeds 2GB limit for HAP_mmap", (uint32_t) size); - abort(); // can't do much else at this point - } - - void *va = HAP_mmap(NULL, size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0); -#endif - if (va == (void*)-1) { - FARF(ERROR, "mmap failed : va %p fd %u size %u", va, fd, (uint32_t) size); + void *va = htp_mmap(fd, size); + if (va == NULL) { + FARF(ERROR, "mmap failed : fd %u size %u", fd, (uint32_t) size); return AEE_EFAILED; } @@ -212,11 +231,7 @@ AEEResult htp_iface_munmap(remote_handle64 handle, uint32 fd) { struct htp_mmap *m = &ctx->mmap[i]; if (fd < 0 || m->fd == fd) { FARF(HIGH, "unmmap : base %p fd %u size %u", (void*) m->base, m->fd, (uint32_t) m->size); -#if __HVX_ARCH__ > 73 - HAP_munmap2((void *) m->base, m->size); -#else - HAP_munmap((void *) m->base, m->size); -#endif + htp_munmap((void *) m->base, m->size); m->size = 0; m->base = NULL; m->fd = -1; @@ -228,7 +243,7 @@ AEEResult htp_iface_munmap(remote_handle64 handle, uint32 fd) { static void vtcm_acquire(struct htp_context * ctx) { if (!ctx->vtcm_valid) { - int err = HAP_compute_res_acquire_cached(ctx->vtcm_rctx, 1000000u); + int err = HAP_compute_res_acquire_cached(ctx->vtcm_rctx, 10000000u); if (err != 0) { FARF(ERROR, "ggml-hex: failed to acquire VTCM: 0x%08x", (unsigned)err); abort(); @@ -692,8 +707,45 @@ static inline void profile_stop(uint32_t mode, struct profile_data * d) { } } +static int op_fence(struct htp_ops_context * octx) { + struct htp_context *ctx = octx->ctx; + struct htp_thread_trace * tr = &ctx->trace[0]; + const uint32_t seq = (uint32_t) octx->op_params[0]; + + htp_trace_event_start(tr, HTP_TRACE_EVT_FENCE, (uint16_t) seq); + + const struct htp_tensor * sync = octx->src[0]; + atomic_uint * sync_fence = (atomic_uint *) sync->data; + uint64_t spins = 0; + while (1) { + Q6_dccleaninva_A((void *) sync_fence); + asm volatile ("syncht" : : : "memory"); + uint32_t val = atomic_load(&sync_fence[0]); + if ((int32_t)(val - seq) >= 0) { + break; + } + if (++spins > HTP_FENCE_TIMEOUT) { + FARF(ERROR, "ggml-hex: sync-wait TIMEOUT : fence %p spins %llu seq %u\n", sync_fence, spins, seq); + break; + } + hex_pause(); + } + + htp_trace_event_stop(tr, HTP_TRACE_EVT_FENCE, (uint16_t) seq); + + FARF(HIGH, "ggml-hex: sync-done : fence %p spins %llu seq %u\n", sync_fence, spins, seq); + return HTP_STATUS_OK; +} + static int execute_op(struct htp_ops_context * octx) { switch (octx->op) { + case HTP_OP_FENCE: + return op_fence(octx); + + case HTP_OP_ALLREDUCE: + case HTP_OP_ALLREDUCE_ADD: + return op_allreduce(octx); + case HTP_OP_MUL_MAT: case HTP_OP_MUL_MAT_ADD: return op_matmul(octx); @@ -701,11 +753,8 @@ static int execute_op(struct htp_ops_context * octx) { case HTP_OP_MUL_MAT_ID: return op_matmul_id(octx); - case HTP_OP_MUL_MAT_QKV: - return op_matmul_qkv(octx); - - case HTP_OP_MUL_MAT_FFN: - return op_matmul_ffn(octx); + case HTP_OP_MUL_MAT_NX: + return op_matmul_nx(octx); case HTP_OP_MUL: case HTP_OP_ADD: @@ -728,6 +777,8 @@ static int execute_op(struct htp_ops_context * octx) { case HTP_OP_UNARY_NEG: case HTP_OP_UNARY_EXP: case HTP_OP_UNARY_TANH: + case HTP_OP_UNARY_ABS: + case HTP_OP_UNARY_LOG: case HTP_OP_L2_NORM: return op_unary(octx); @@ -818,12 +869,8 @@ static inline bool reuse_buf(struct htp_context *ctx, uint32_t *m_reuse, struct static inline void drop_mmap(struct htp_context *ctx, struct htp_mmap *m) { if (m->size) { - FARF(HIGH, "unmap : fd %u base %p size %u", m->fd, (void*) m->base, (uint32_t) m->size); -#if __HVX_ARCH__ > 73 - HAP_munmap2((void *) m->base, m->size); -#else - HAP_munmap((void *) m->base, m->size); -#endif + FARF(ALWAYS, "unmap : fd %u base %p size %u", m->fd, (void*) m->base, (uint32_t) m->size); + htp_munmap((void *) m->base, m->size); m->size = 0; m->base = 0; m->fd = -1; @@ -837,18 +884,9 @@ static inline void mmap_buf(struct htp_context *ctx, struct htp_buf_desc *b) { for (uint32_t i=0; i < HTP_MAX_MMAPS; i++) { struct htp_mmap *m = &ctx->mmap[i]; if (!m->size) { -#if __HVX_ARCH__ > 73 - void *va = HAP_mmap2(NULL, b->size, HAP_PROT_READ | HAP_PROT_WRITE, 0, b->fd, 0); -#else - if (b->size > HTP_MMAP_MAX_VMEM) { // HAP_mmap has a size limit of 2GB - FARF(ERROR, "mmap failed : size %u exceeds 2GB limit for HAP_mmap", (uint32_t) b->size); - abort(); // can't do much else at this point - } - - void *va = HAP_mmap(NULL, b->size, HAP_PROT_READ | HAP_PROT_WRITE, 0, b->fd, 0); -#endif - if (va == (void*)-1) { - FARF(ERROR, "mmap failed : va %p fd %u size %u", va, b->fd, (uint32_t) b->size); + void *va = htp_mmap(b->fd, b->size); + if (va == NULL) { + FARF(ERROR, "mmap failed : fd %u size %u", b->fd, (uint32_t) b->size); abort(); // can't do much else at this point } @@ -856,10 +894,13 @@ static inline void mmap_buf(struct htp_context *ctx, struct htp_buf_desc *b) { m->fd = b->fd; m->size = b->size; - FARF(HIGH, "mmap : fd %u base %p size %u", m->fd, (void*) m->base, (uint32_t) m->size); + FARF(ALWAYS, "mmap : fd %u base %p size %u", m->fd, (void*) m->base, (uint32_t) m->size); return; } } + + FARF(ERROR, "mmap failed : exceeded mapping capacity limit of %u", HTP_MAX_MMAPS); + abort(); } static void prep_op_bufs(struct htp_context *ctx, struct htp_buf_desc *bufs, uint32_t n_bufs) { @@ -1081,6 +1122,7 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r rsp.usecs = batch_prof.usecs; rsp.cycles_start = batch_prof.cycles_start; rsp.cycles_stop = batch_prof.cycles_stop; + rsp.seq = req->seq; if (ctx->profiler == HTP_PROF_TRACE) { for (int t = 0; t <= HTP_MAX_NTHREADS; t++) { diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index 9d385469ae9..a6adc0e61fa 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -64,6 +64,7 @@ typedef struct { struct htp_mm_context { const char * type; struct htp_ops_context * octx; + const struct htp_tensor * act; void (*vec_dot_1x1)(const uint32_t n, float * restrict s0, const void * restrict vx0, @@ -478,7 +479,7 @@ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void \ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ DOT_2X1(ne10, dst_ptr, w_tile, src1_col, valid_rows, NULL); \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ \ if (push_ct < ct_end) { \ dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src0_row + push_ct * tile_row_stride), \ @@ -502,150 +503,67 @@ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void } \ } -#define MATMUL_QKV_2D_REPACKED_IMPL(SUFFIX, TILE_SIZE, DOT_2X2, DOT_2X1) \ -static void hvx_mm_qkv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ +#define MATMUL_NX_2D_REPACKED_IMPL(SUFFIX, TILE_SIZE, DOT_2X2, DOT_2X1) \ +static void hvx_mm_nx_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ struct htp_mm_context * mmctx = data; \ struct htp_ops_context * octx = mmctx->octx; \ + const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \ + const uint32_t n_weights = kparams->n_weights; \ \ - const struct htp_tensor * restrict src0 = octx->src[0]; /* Wk */ \ - const struct htp_tensor * restrict src1 = octx->src[1]; /* x */ \ - const struct htp_tensor * restrict src2 = octx->src[2]; /* Wv */ \ - const struct htp_tensor * restrict src3 = octx->src[3]; /* Wq */ \ - const struct htp_tensor * restrict dst_k = octx->dsts[0]; \ - const struct htp_tensor * restrict dst_v = octx->dsts[1]; \ - const struct htp_tensor * restrict dst_q = octx->dsts[2]; \ - \ - const uint32_t ne00 = src0->ne[0]; \ - const uint32_t ne10 = src1->ne[0]; \ - const uint32_t src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3]; \ - \ - const size_t dst_k_row_size = dst_k->nb[1]; /* K and V share output width */ \ - const size_t dst_q_row_size = dst_q->nb[1]; /* Q may be wider (GQA) */ \ + const struct htp_tensor * restrict act = octx->src[n_weights]; /* x */ \ + const uint32_t ne10 = act->ne[0]; \ + const uint32_t src1_nrows = act->ne[1] * act->ne[2] * act->ne[3]; \ const size_t src1_stride = mmctx->vtcm_src1_stride; \ \ - uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; \ - uint8_t * restrict vtcm_src2_ptr = mmctx->vtcm_src2 + mmctx->vtcm_src2_size_per_thread * ith; \ - uint8_t * restrict vtcm_src3_ptr = mmctx->vtcm_src3 + mmctx->vtcm_src3_size_per_thread * ith; \ - uint8_t * restrict src1_data = mmctx->vtcm_src1; \ + uint8_t * restrict vtcm_weight_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; \ + uint8_t * restrict src1_data = mmctx->vtcm_src1; \ \ struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ - \ - const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \ const uint32_t n_prefetch = kparams->n_prefetch; \ assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); \ \ - const uint8_t * restrict src0_row = (const uint8_t *) src0->data; \ - const uint8_t * restrict src2_row = (const uint8_t *) src2->data; \ - const uint8_t * restrict src3_row = (const uint8_t *) src3->data; \ - \ const uint32_t tile_size = TILE_SIZE; \ const uint32_t aligned_tile_size = hex_align_up(tile_size, 128); \ - \ - uint32_t n_k_tiles_w = ne00 / 32; \ uint32_t n_k_tiles_a = ne10 / 32; \ - uint32_t tile_row_stride = n_k_tiles_w * tile_size; \ uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; \ \ dma_queue * dma_queue = octx->ctx->dma[ith]; \ \ - /* 1. Process K and V together */ \ - const uint32_t src0_nrows_kv = src0->ne[1] * src0->ne[2] * src0->ne[3]; /* src0 is Wk */ \ - uint32_t src0_nrows_per_thread_kv = (src0_nrows_kv + nth - 1) / nth; \ - src0_nrows_per_thread_kv = hex_round_up(src0_nrows_per_thread_kv, 32); \ - \ - const uint32_t start_row_kv = src0_nrows_per_thread_kv * ith; \ - const uint32_t end_row_kv = MIN(start_row_kv + src0_nrows_per_thread_kv, src0_nrows_kv); \ - \ - uint32_t ct_start_kv = start_row_kv / 32; \ - uint32_t ct_end_kv = (end_row_kv + 31) / 32; \ - \ - uint32_t push_ct = ct_start_kv; \ - if (start_row_kv < end_row_kv) { \ - for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end_kv; d++, push_ct++) { \ - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, \ - src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + d * tile_row_transfer_size_aligned, \ - src2_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - } \ - } \ - \ hvx_mm_run_quant_task(mmctx, ith); \ \ - if (start_row_kv < end_row_kv) { \ - \ - for (uint32_t ct = ct_start_kv; ct < ct_end_kv; ct++) { \ - const uint8_t * w_tile_k = dma_queue_pop(dma_queue).dst; \ - const uint8_t * w_tile_v = dma_queue_pop(dma_queue).dst; \ - \ - int valid_rows = (int)src0->ne[1] - (int)(ct * 32); \ - valid_rows = MIN(32, MAX(0, valid_rows)); \ - \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ith); \ - uint32_t ir1 = 0; \ - for (; ir1 + 1 < src1_nrows; ir1 += 2) { \ - const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (ir1+0) * src1_stride); \ - const uint8_t * restrict src1_col1 = (const uint8_t *) (src1_data + (ir1+1) * src1_stride); \ - \ - float * restrict dst_row0_k = (float *) (dst_k->data + ((ir1+0) * dst_k_row_size)); \ - float * restrict dst_row1_k = (float *) (dst_k->data + ((ir1+1) * dst_k_row_size)); \ - float * dst_ptr0_k = &dst_row0_k[ct * 32]; \ - float * dst_ptr1_k = &dst_row1_k[ct * 32]; \ - \ - float * restrict dst_row0_v = (float *) (dst_v->data + ((ir1+0) * dst_k_row_size)); \ - float * restrict dst_row1_v = (float *) (dst_v->data + ((ir1+1) * dst_k_row_size)); \ - float * dst_ptr0_v = &dst_row0_v[ct * 32]; \ - float * dst_ptr1_v = &dst_row1_v[ct * 32]; \ - \ - DOT_2X2(ne10, dst_ptr0_k, dst_ptr1_k, w_tile_k, src1_col0, src1_col1, valid_rows, NULL, NULL); \ - DOT_2X2(ne10, dst_ptr0_v, dst_ptr1_v, w_tile_v, src1_col0, src1_col1, valid_rows, NULL, NULL); \ - } \ + for (uint32_t widx = 0; widx < n_weights; widx++) { \ + const struct htp_tensor * restrict src_w = octx->src[widx]; \ + const struct htp_tensor * restrict dst = octx->dsts[widx]; \ + if (!src_w || !dst) continue; \ \ - for (; ir1 < src1_nrows; ++ir1) { \ - const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); \ - \ - float * restrict dst_row_k = (float *) (dst_k->data + (ir1 * dst_k_row_size)); \ - float * dst_ptr_k = &dst_row_k[ct * 32]; \ - \ - float * restrict dst_row_v = (float *) (dst_v->data + (ir1 * dst_k_row_size)); \ - float * dst_ptr_v = &dst_row_v[ct * 32]; \ - \ - DOT_2X1(ne10, dst_ptr_k, w_tile_k, src1_col, valid_rows, NULL); \ - DOT_2X1(ne10, dst_ptr_v, w_tile_v, src1_col, valid_rows, NULL); \ - } \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ith); \ + const uint32_t ne00 = src_w->ne[0]; \ + const uint32_t ne01 = src_w->ne[1]; \ + const size_t dst_row_size = dst->nb[1]; \ + const uint8_t * restrict src_w_row = (const uint8_t *) src_w->data; \ \ - if (push_ct < ct_end_kv) { \ - dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile_k, src0_row + push_ct * tile_row_stride), \ - aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile_v, src2_row + push_ct * tile_row_stride), \ - aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - push_ct++; \ - } \ - } \ - } \ + uint32_t n_k_tiles_w = ne00 / 32; \ + uint32_t tile_row_stride = n_k_tiles_w * tile_size; \ \ - /* 2. Process Q separately */ \ - const uint32_t src0_nrows_q = src3->ne[1] * src3->ne[2] * src3->ne[3]; /* src3 is Wq */ \ - uint32_t src0_nrows_per_thread_q = (src0_nrows_q + nth - 1) / nth; \ - src0_nrows_per_thread_q = hex_round_up(src0_nrows_per_thread_q, 32); \ + const uint32_t src0_nrows = ne01 * src_w->ne[2] * src_w->ne[3]; \ + uint32_t src0_nrows_per_thread = (src0_nrows + nth - 1) / nth; \ + src0_nrows_per_thread = hex_round_up(src0_nrows_per_thread, 32); \ \ - const uint32_t start_row_q = src0_nrows_per_thread_q * ith; \ - const uint32_t end_row_q = MIN(start_row_q + src0_nrows_per_thread_q, src0_nrows_q); \ + const uint32_t start_row = src0_nrows_per_thread * ith; \ + const uint32_t end_row = MIN(start_row + src0_nrows_per_thread, src0_nrows); \ + if (start_row >= end_row) continue; \ \ - if (start_row_q < end_row_q) { \ - uint32_t ct_start_q = start_row_q / 32; \ - uint32_t ct_end_q = (end_row_q + 31) / 32; \ + uint32_t ct_start = start_row / 32; \ + uint32_t ct_end = (end_row + 31) / 32; \ \ - uint32_t push_ct = ct_start_q; \ - for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end_q; d++, push_ct++) { \ - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src3_ptr + d * tile_row_transfer_size_aligned, \ - src3_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ + uint32_t push_ct = ct_start; \ + for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { \ + dma_queue_push(dma_queue, dma_make_ptr(vtcm_weight_ptr + d * tile_row_transfer_size_aligned, \ + src_w_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ } \ \ - for (uint32_t ct = ct_start_q; ct < ct_end_q; ct++) { \ - const uint8_t * w_tile_q = dma_queue_pop(dma_queue).dst; \ - \ - int valid_rows = (int)src3->ne[1] - (int)(ct * 32); \ + for (uint32_t ct = ct_start; ct < ct_end; ct++) { \ + const uint8_t * w_tile = dma_queue_pop(dma_queue).dst; \ + int valid_rows = (int)ne01 - (int)(ct * 32); \ valid_rows = MIN(32, MAX(0, valid_rows)); \ \ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ @@ -654,26 +572,24 @@ static void hvx_mm_qkv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (ir1+0) * src1_stride); \ const uint8_t * restrict src1_col1 = (const uint8_t *) (src1_data + (ir1+1) * src1_stride); \ \ - float * restrict dst_row0_q = (float *) (dst_q->data + ((ir1+0) * dst_q_row_size)); \ - float * restrict dst_row1_q = (float *) (dst_q->data + ((ir1+1) * dst_q_row_size)); \ - float * dst_ptr0_q = &dst_row0_q[ct * 32]; \ - float * dst_ptr1_q = &dst_row1_q[ct * 32]; \ + float * restrict dst_row0 = (float *) (dst->data + ((ir1+0) * dst_row_size)); \ + float * restrict dst_row1 = (float *) (dst->data + ((ir1+1) * dst_row_size)); \ + float * dst_ptr0 = &dst_row0[ct * 32]; \ + float * dst_ptr1 = &dst_row1[ct * 32]; \ \ - DOT_2X2(ne10, dst_ptr0_q, dst_ptr1_q, w_tile_q, src1_col0, src1_col1, valid_rows, NULL, NULL); \ + DOT_2X2(ne10, dst_ptr0, dst_ptr1, w_tile, src1_col0, src1_col1, valid_rows, NULL, NULL); \ } \ \ for (; ir1 < src1_nrows; ++ir1) { \ const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); \ - \ - float * restrict dst_row_q = (float *) (dst_q->data + (ir1 * dst_q_row_size)); \ - float * dst_ptr_q = &dst_row_q[ct * 32]; \ - \ - DOT_2X1(ne10, dst_ptr_q, w_tile_q, src1_col, valid_rows, NULL); \ + float * restrict dst_row = (float *) (dst->data + (ir1 * dst_row_size)); \ + float * dst_ptr = &dst_row[ct * 32]; \ + DOT_2X1(ne10, dst_ptr, w_tile, src1_col, valid_rows, NULL); \ } \ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ \ - if (push_ct < ct_end_q) { \ - dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile_q, src3_row + push_ct * tile_row_stride), \ + if (push_ct < ct_end) { \ + dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src_w_row + push_ct * tile_row_stride), \ aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ push_ct++; \ } \ @@ -681,121 +597,6 @@ static void hvx_mm_qkv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, } \ } -#define MATMUL_FFN_2D_REPACKED_IMPL(SUFFIX, TILE_SIZE, DOT_2X2, DOT_2X1) \ -static void hvx_mm_ffn_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ - struct htp_mm_context * mmctx = data; \ - struct htp_ops_context * octx = mmctx->octx; \ - \ - const struct htp_tensor * restrict src0 = octx->src[0]; /* Wgate */ \ - const struct htp_tensor * restrict src1 = octx->src[1]; /* y */ \ - const struct htp_tensor * restrict src2 = octx->src[2]; /* Wup */ \ - const struct htp_tensor * restrict dst_gate = octx->dsts[0]; \ - const struct htp_tensor * restrict dst_up = octx->dsts[1]; \ - \ - const uint32_t ne00 = src0->ne[0]; \ - const uint32_t ne01 = src0->ne[1]; \ - const uint32_t ne10 = src1->ne[0]; \ - const uint32_t src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3]; \ - \ - const size_t dst_row_size = dst_gate->nb[1]; \ - const size_t src1_stride = mmctx->vtcm_src1_stride; \ - \ - uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; \ - uint8_t * restrict vtcm_src2_ptr = mmctx->vtcm_src2 + mmctx->vtcm_src2_size_per_thread * ith; \ - uint8_t * restrict src1_data = mmctx->vtcm_src1; \ - \ - struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ - \ - const uint8_t * restrict src0_row = (const uint8_t *) src0->data; \ - const uint8_t * restrict src2_row = (const uint8_t *) src2->data; \ - \ - const uint32_t tile_size = TILE_SIZE; \ - const uint32_t aligned_tile_size = hex_align_up(tile_size, 128); \ - \ - const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \ - const uint32_t n_prefetch = kparams->n_prefetch; \ - assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); \ - \ - uint32_t n_k_tiles_w = ne00 / 32; \ - uint32_t n_k_tiles_a = ne10 / 32; \ - uint32_t tile_row_stride = n_k_tiles_w * tile_size; \ - uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; \ - dma_queue * dma_queue = octx->ctx->dma[ith]; \ - \ - const uint32_t src0_nrows = ne01 * src0->ne[2] * src0->ne[3]; \ - const uint32_t src0_start_row = mmctx->src0_nrows_per_thread * ith; \ - const uint32_t src0_end_row = MIN(src0_start_row + mmctx->src0_nrows_per_thread, src0_nrows); \ - \ - uint32_t ct_start = src0_start_row / 32; \ - uint32_t ct_end = (src0_end_row + 31) / 32; \ - \ - uint32_t push_ct = ct_start; \ - if (src0_start_row < src0_end_row) { \ - for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { \ - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, \ - src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + d * tile_row_transfer_size_aligned, \ - src2_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - } \ - } \ - \ - hvx_mm_run_quant_task(mmctx, ith); \ - \ - if (src0_start_row >= src0_end_row) { \ - return; \ - } \ - \ - for (uint32_t ct = ct_start; ct < ct_end; ct++) { \ - const uint8_t * w_tile_gate = dma_queue_pop(dma_queue).dst; \ - const uint8_t * w_tile_up = dma_queue_pop(dma_queue).dst; \ - \ - int valid_rows = (int)ne01 - (int)(ct * 32); \ - valid_rows = MIN(32, MAX(0, valid_rows)); \ - \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ - uint32_t ir1 = 0; \ - for (; ir1 + 1 < src1_nrows; ir1 += 2) { \ - const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (ir1+0) * src1_stride); \ - const uint8_t * restrict src1_col1 = (const uint8_t *) (src1_data + (ir1+1) * src1_stride); \ - \ - float * restrict dst_row0_gate = (float *) (dst_gate->data + ((ir1+0) * dst_row_size)); \ - float * restrict dst_row1_gate = (float *) (dst_gate->data + ((ir1+1) * dst_row_size)); \ - float * dst_ptr0_gate = &dst_row0_gate[ct * 32]; \ - float * dst_ptr1_gate = &dst_row1_gate[ct * 32]; \ - \ - float * restrict dst_row0_up = (float *) (dst_up->data + ((ir1+0) * dst_row_size)); \ - float * restrict dst_row1_up = (float *) (dst_up->data + ((ir1+1) * dst_row_size)); \ - float * dst_ptr0_up = &dst_row0_up[ct * 32]; \ - float * dst_ptr1_up = &dst_row1_up[ct * 32]; \ - \ - DOT_2X2(ne10, dst_ptr0_gate, dst_ptr1_gate, w_tile_gate, src1_col0, src1_col1, valid_rows, NULL, NULL); \ - DOT_2X2(ne10, dst_ptr0_up, dst_ptr1_up, w_tile_up, src1_col0, src1_col1, valid_rows, NULL, NULL); \ - } \ - \ - for (; ir1 < src1_nrows; ++ir1) { \ - const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); \ - \ - float * restrict dst_row_gate = (float *) (dst_gate->data + (ir1 * dst_row_size)); \ - float * dst_ptr_gate = &dst_row_gate[ct * 32]; \ - \ - float * restrict dst_row_up = (float *) (dst_up->data + (ir1 * dst_row_size)); \ - float * dst_ptr_up = &dst_row_up[ct * 32]; \ - \ - DOT_2X1(ne10, dst_ptr_gate, w_tile_gate, src1_col, valid_rows, NULL); \ - DOT_2X1(ne10, dst_ptr_up, w_tile_up, src1_col, valid_rows, NULL); \ - } \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ - \ - if (push_ct < ct_end) { \ - dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile_gate, src0_row + push_ct * tile_row_stride), \ - aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile_up, src2_row + push_ct * tile_row_stride), \ - aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - push_ct++; \ - } \ - } \ -} - MATMUL_2D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x2, tiled_vec_dot_q4_0_32x1) MATMUL_2D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x2, tiled_vec_dot_q4_1_32x1) MATMUL_2D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x2, tiled_vec_dot_q8_0_32x1) @@ -812,7 +613,7 @@ MATMUL_2D_REPACKED_IMPL(mxfp4_flat, 544, flat_vec_dot_mxfp4_32x2, flat_vec_dot static void name(unsigned int nth, unsigned int ith, void * data) { \ struct htp_mm_context * mmctx = data; \ struct htp_ops_context * octx = mmctx->octx; \ - const struct htp_tensor * src = octx->src[1]; \ + const struct htp_tensor * src = mmctx->act; \ const uint32_t ne0 = src->ne[0]; \ const uint32_t ne1 = src->ne[1]; \ const uint32_t ne2 = src->ne[2]; \ @@ -854,7 +655,7 @@ static void quantize_f32_q8_0_tiled_block(unsigned int nth, unsigned int ith, vo struct htp_thread_trace * tr = &octx->ctx->trace[ith]; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, mmctx->quant_ib_first[ith]); - const struct htp_tensor * src = octx->src[1]; + const struct htp_tensor * src = mmctx->act; quantize_f32_q8_0_tiled_block_kernel( (const float *) src->data, @@ -878,7 +679,7 @@ static void quantize_f32_q8_1_tiled_block(unsigned int nth, unsigned int ith, vo struct htp_thread_trace * tr = &octx->ctx->trace[ith]; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, mmctx->quant_ib_first[ith]); - const struct htp_tensor * src = octx->src[1]; + const struct htp_tensor * src = mmctx->act; quantize_f32_q8_1_tiled_block_kernel( (const float *) src->data, @@ -909,30 +710,17 @@ MATVEC_2D_REPACKED_IMPL(iq4nl_flat, 576, flat_vec_dot_iq4nl_32x1) MATVEC_2D_REPACKED_IMPL(mxfp4_flat, 544, flat_vec_dot_mxfp4_32x1) -MATMUL_QKV_2D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x2, tiled_vec_dot_q4_0_32x1) -MATMUL_QKV_2D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x2, tiled_vec_dot_q4_1_32x1) -MATMUL_QKV_2D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x2, tiled_vec_dot_q8_0_32x1) -MATMUL_QKV_2D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x2, tiled_vec_dot_iq4nl_32x1) -MATMUL_QKV_2D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x2, tiled_vec_dot_mxfp4_32x1) +MATMUL_NX_2D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x2, tiled_vec_dot_q4_0_32x1) +MATMUL_NX_2D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x2, tiled_vec_dot_q4_1_32x1) +MATMUL_NX_2D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x2, tiled_vec_dot_q8_0_32x1) +MATMUL_NX_2D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x2, tiled_vec_dot_iq4nl_32x1) +MATMUL_NX_2D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x2, tiled_vec_dot_mxfp4_32x1) -MATMUL_QKV_2D_REPACKED_IMPL(q4_0_flat, 576, flat_vec_dot_q4_0_32x2, flat_vec_dot_q4_0_32x1) -MATMUL_QKV_2D_REPACKED_IMPL(q4_1_flat, 640, flat_vec_dot_q4_1_32x2, flat_vec_dot_q4_1_32x1) -MATMUL_QKV_2D_REPACKED_IMPL(q8_0_flat, 1088, flat_vec_dot_q8_0_32x2, flat_vec_dot_q8_0_32x1) -MATMUL_QKV_2D_REPACKED_IMPL(iq4nl_flat, 576, flat_vec_dot_iq4nl_32x2, flat_vec_dot_iq4nl_32x1) -MATMUL_QKV_2D_REPACKED_IMPL(mxfp4_flat, 544, flat_vec_dot_mxfp4_32x2, flat_vec_dot_mxfp4_32x1) - - -MATMUL_FFN_2D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x2, tiled_vec_dot_q4_0_32x1) -MATMUL_FFN_2D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x2, tiled_vec_dot_q4_1_32x1) -MATMUL_FFN_2D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x2, tiled_vec_dot_q8_0_32x1) -MATMUL_FFN_2D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x2, tiled_vec_dot_iq4nl_32x1) -MATMUL_FFN_2D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x2, tiled_vec_dot_mxfp4_32x1) - -MATMUL_FFN_2D_REPACKED_IMPL(q4_0_flat, 576, flat_vec_dot_q4_0_32x2, flat_vec_dot_q4_0_32x1) -MATMUL_FFN_2D_REPACKED_IMPL(q4_1_flat, 640, flat_vec_dot_q4_1_32x2, flat_vec_dot_q4_1_32x1) -MATMUL_FFN_2D_REPACKED_IMPL(q8_0_flat, 1088, flat_vec_dot_q8_0_32x2, flat_vec_dot_q8_0_32x1) -MATMUL_FFN_2D_REPACKED_IMPL(iq4nl_flat, 576, flat_vec_dot_iq4nl_32x2, flat_vec_dot_iq4nl_32x1) -MATMUL_FFN_2D_REPACKED_IMPL(mxfp4_flat, 544, flat_vec_dot_mxfp4_32x2, flat_vec_dot_mxfp4_32x1) +MATMUL_NX_2D_REPACKED_IMPL(q4_0_flat, 576, flat_vec_dot_q4_0_32x2, flat_vec_dot_q4_0_32x1) +MATMUL_NX_2D_REPACKED_IMPL(q4_1_flat, 640, flat_vec_dot_q4_1_32x2, flat_vec_dot_q4_1_32x1) +MATMUL_NX_2D_REPACKED_IMPL(q8_0_flat, 1088, flat_vec_dot_q8_0_32x2, flat_vec_dot_q8_0_32x1) +MATMUL_NX_2D_REPACKED_IMPL(iq4nl_flat, 576, flat_vec_dot_iq4nl_32x2, flat_vec_dot_iq4nl_32x1) +MATMUL_NX_2D_REPACKED_IMPL(mxfp4_flat, 544, flat_vec_dot_mxfp4_32x2, flat_vec_dot_mxfp4_32x1) static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) { htp_matmul_preamble; @@ -1353,6 +1141,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { struct htp_mm_context mmctx_struct = {0}; struct htp_mm_context * mmctx = &mmctx_struct; mmctx->octx = octx; + mmctx->act = src1; const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; @@ -1528,7 +1317,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { struct htp_mm_hvx_vtcm_layout L; htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads, - dst_row_size, src0_row_size, src1_row_size, src2 ? src2->nb[1] : 0, kparams->n_prefetch, false, false, false); + dst_row_size, src0_row_size, src1_row_size, src2 ? src2->nb[1] : 0, kparams->n_prefetch, false, false); if (kparams->kernel_type == HTP_MM_KERNEL_HVX_F16_F16_VTCM || kparams->kernel_type == HTP_MM_KERNEL_HVX_F32_F32_VTCM || @@ -1587,297 +1376,97 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { return HTP_STATUS_OK; } -static void hvx_mm_qkv_2d(unsigned int nth, unsigned int ith, void * data) { +static void hvx_mm_nx_2d(unsigned int nth, unsigned int ith, void * data) { struct htp_mm_context * mmctx = data; struct htp_ops_context * octx = mmctx->octx; + const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; + const uint32_t n_weights = kparams->n_weights; - const struct htp_tensor * restrict src0 = octx->src[0]; // Wk - const struct htp_tensor * restrict src1 = octx->src[1]; // x - const struct htp_tensor * restrict src2 = octx->src[2]; // Wv - const struct htp_tensor * restrict src3 = octx->src[3]; // Wq - const struct htp_tensor * restrict dst_k = octx->dsts[0]; - const struct htp_tensor * restrict dst_v = octx->dsts[1]; - const struct htp_tensor * restrict dst_q = octx->dsts[2]; - - const uint32_t ne00 = src0->ne[0]; - const uint32_t ne01 = src0->ne[1]; - const uint32_t ne02 = src0->ne[2]; - const uint32_t ne03 = src0->ne[3]; - - const uint32_t ne11 = src1->ne[1]; - const uint32_t ne12 = src1->ne[2]; - const uint32_t ne13 = src1->ne[3]; - - const uint32_t src0_nrows = ne01 * ne02 * ne03; - const uint32_t src1_nrows = ne11 * ne12 * ne13; - - const uint32_t src0_nrows_per_thread = mmctx->src0_nrows_per_thread; - const uint32_t src0_start_row = src0_nrows_per_thread * ith; - const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows); - const uint32_t src0_end_row_x2 = src0_start_row + ((src0_end_row - src0_start_row) & ~1U); - - const size_t dst_k_row_size = dst_k->nb[1]; // K and V share output width - const size_t dst_q_row_size = dst_q->nb[1]; // Q may be wider (GQA) - const size_t src0_row_size = src0->nb[1]; - const size_t src2_row_size = src2->nb[1]; - const size_t src3_row_size = src3->nb[1]; - - const size_t src0_stride = mmctx->vtcm_src0_stride; - const size_t src2_stride = mmctx->vtcm_src2_stride; - const size_t src3_stride = mmctx->vtcm_src3_stride; + const struct htp_tensor * restrict act = octx->src[n_weights]; + const uint32_t src1_nrows = act->ne[1] * act->ne[2] * act->ne[3]; const size_t src1_stride = mmctx->vtcm_src1_stride; uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; - uint8_t * restrict vtcm_src2_ptr = mmctx->vtcm_src2 + mmctx->vtcm_src2_size_per_thread * ith; - uint8_t * restrict vtcm_src3_ptr = mmctx->vtcm_src3 + mmctx->vtcm_src3_size_per_thread * ith; - uint8_t * restrict src1_data = mmctx->vtcm_src1; + uint8_t * restrict src1_data = mmctx->vtcm_src1; dma_queue * dma_queue = octx->ctx->dma[ith]; - - const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; const uint32_t n_prefetch = kparams->n_prefetch; assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); const uint32_t prefetch_mask = n_prefetch - 1; - const uint8_t * restrict src0_row = (const uint8_t *) src0->data; - const uint8_t * restrict src2_row = (const uint8_t *) src2->data; - const uint8_t * restrict src3_row = (const uint8_t *) src3->data; - - // Prefill spad with src0, src2, src3 rows - if (src0_start_row < src0_end_row) { - for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { - const int is0 = (ir0 - src0_start_row); - if (is0 >= (int)n_prefetch) { - break; - } - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), - src0_stride, src0_row_size, src0_row_size, 2); - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + is0 * src2_stride, src2_row + ir0 * src2_row_size), - src2_stride, src2_row_size, src2_row_size, 2); - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src3_ptr + is0 * src3_stride, src3_row + ir0 * src3_row_size), - src3_stride, src3_row_size, src3_row_size, 2); - } - } + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; hvx_mm_run_quant_task(mmctx, ith); - if (src0_start_row >= src0_end_row) { - return; - } - - // Process rows - for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { - const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; - const uint8_t * ss2 = dma_queue_pop(dma_queue).dst; - const uint8_t * ss3 = dma_queue_pop(dma_queue).dst; - - // Process src1 columns in pairs (2×2 tiling) - uint32_t ir1 = 0; - for (; ir1 + 1 < src1_nrows; ir1 += 2) { - const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (ir1+0) * src1_stride); - const uint8_t * restrict src1_col1 = (const uint8_t *) (src1_data + (ir1+1) * src1_stride); - - float * restrict dst_row0_k = (float *) (dst_k->data + ((ir1+0) * dst_k_row_size)); - float * restrict dst_row1_k = (float *) (dst_k->data + ((ir1+1) * dst_k_row_size)); - mmctx->vec_dot_2x2(ne00, &dst_row0_k[ir0], &dst_row1_k[ir0], ss0, ss0 + src0_stride, src1_col0, src1_col1); - - float * restrict dst_row0_v = (float *) (dst_v->data + ((ir1+0) * dst_k_row_size)); - float * restrict dst_row1_v = (float *) (dst_v->data + ((ir1+1) * dst_k_row_size)); - mmctx->vec_dot_2x2(ne00, &dst_row0_v[ir0], &dst_row1_v[ir0], ss2, ss2 + src2_stride, src1_col0, src1_col1); - - float * restrict dst_row0_q = (float *) (dst_q->data + ((ir1+0) * dst_q_row_size)); - float * restrict dst_row1_q = (float *) (dst_q->data + ((ir1+1) * dst_q_row_size)); - mmctx->vec_dot_2x2(ne00, &dst_row0_q[ir0], &dst_row1_q[ir0], ss3, ss3 + src3_stride, src1_col0, src1_col1); - } - - // Handle remaining src1 rows (fallback to 2×1) - for (; ir1 < src1_nrows; ++ir1) { - const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); - - float * restrict dst_row_k = (float *) (dst_k->data + (ir1 * dst_k_row_size)); - mmctx->vec_dot_2x1(ne00, &dst_row_k[ir0], ss0, ss0 + src0_stride, src1_col); - - float * restrict dst_row_v = (float *) (dst_v->data + (ir1 * dst_k_row_size)); - mmctx->vec_dot_2x1(ne00, &dst_row_v[ir0], ss2, ss2 + src2_stride, src1_col); - - float * restrict dst_row_q = (float *) (dst_q->data + (ir1 * dst_q_row_size)); - mmctx->vec_dot_2x1(ne00, &dst_row_q[ir0], ss3, ss3 + src3_stride, src1_col); - } - - // Prefetch next (n + vtcm_nrows) rows - const int pr0 = (ir0 + n_prefetch); - const int is0 = (pr0 - src0_start_row) & prefetch_mask; - if (pr0 < src0_end_row_x2) { - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + pr0 * src0_row_size), - src0_stride, src0_row_size, src0_row_size, 2); - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + is0 * src2_stride, src2_row + pr0 * src2_row_size), - src2_stride, src2_row_size, src2_row_size, 2); - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src3_ptr + is0 * src3_stride, src3_row + pr0 * src3_row_size), - src3_stride, src3_row_size, src3_row_size, 2); - } - } - - // Process last row (if any) - if (src0_end_row != src0_end_row_x2) { - uint32_t ir0 = src0_end_row_x2; - const int is0 = (ir0 - src0_start_row) & prefetch_mask; - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), - src0_stride, src0_row_size, src0_row_size, 1); - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + is0 * src2_stride, src2_row + ir0 * src2_row_size), - src2_stride, src2_row_size, src2_row_size, 1); - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src3_ptr + is0 * src3_stride, src3_row + ir0 * src3_row_size), - src3_stride, src3_row_size, src3_row_size, 1); - - const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; - const uint8_t * ss2 = dma_queue_pop(dma_queue).dst; - const uint8_t * ss3 = dma_queue_pop(dma_queue).dst; - - for (uint32_t ir1 = 0; ir1 < src1_nrows; ++ir1) { - const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); - - float * restrict dst_row_k = (float *) (dst_k->data + (ir1 * dst_k_row_size)); - mmctx->vec_dot_1x1(ne00, &dst_row_k[ir0], ss0, src1_col); + for (uint32_t widx = 0; widx < n_weights; widx++) { + const struct htp_tensor * restrict src_w = octx->src[widx]; + const struct htp_tensor * restrict dst = octx->dsts[widx]; + if (!src_w || !dst) continue; - float * restrict dst_row_v = (float *) (dst_v->data + (ir1 * dst_k_row_size)); - mmctx->vec_dot_1x1(ne00, &dst_row_v[ir0], ss2, src1_col); + const uint32_t ne00 = src_w->ne[0]; + const uint32_t ne01 = src_w->ne[1]; + const uint32_t src0_nrows = ne01 * src_w->ne[2] * src_w->ne[3]; - float * restrict dst_row_q = (float *) (dst_q->data + (ir1 * dst_q_row_size)); - mmctx->vec_dot_1x1(ne00, &dst_row_q[ir0], ss3, src1_col); - } - } -} - -static void hvx_mm_ffn_2d(unsigned int nth, unsigned int ith, void * data) { - struct htp_mm_context * mmctx = data; - struct htp_ops_context * octx = mmctx->octx; + uint32_t src0_nrows_per_thread = (src0_nrows + nth - 1) / nth; + src0_nrows_per_thread += (src0_nrows_per_thread & 1); - const struct htp_tensor * restrict src0 = octx->src[0]; // Wgate - const struct htp_tensor * restrict src1 = octx->src[1]; // y - const struct htp_tensor * restrict src2 = octx->src[2]; // Wup - const struct htp_tensor * restrict dst_gate = octx->dsts[0]; - const struct htp_tensor * restrict dst_up = octx->dsts[1]; + const uint32_t src0_start_row = src0_nrows_per_thread * ith; + const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows); + const uint32_t src0_end_row_x2 = src0_start_row + ((src0_end_row - src0_start_row) & ~1U); + if (src0_start_row >= src0_end_row) continue; - const uint32_t ne00 = src0->ne[0]; - const uint32_t ne01 = src0->ne[1]; - const uint32_t ne02 = src0->ne[2]; - const uint32_t ne03 = src0->ne[3]; + const size_t dst_row_size = dst->nb[1]; + const size_t src0_row_size = src_w->nb[1]; + const size_t src0_stride = hex_round_up(src0_row_size, 128); - const uint32_t ne11 = src1->ne[1]; - const uint32_t ne12 = src1->ne[2]; - const uint32_t ne13 = src1->ne[3]; + const uint8_t * restrict src0_row = (const uint8_t *) src_w->data; - const uint32_t src0_nrows = ne01 * ne02 * ne03; - const uint32_t src1_nrows = ne11 * ne12 * ne13; - - const uint32_t src0_nrows_per_thread = mmctx->src0_nrows_per_thread; - const uint32_t src0_start_row = src0_nrows_per_thread * ith; - const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows); - const uint32_t src0_end_row_x2 = src0_start_row + ((src0_end_row - src0_start_row) & ~1U); - - const size_t dst_row_size = dst_gate->nb[1]; - const size_t src0_row_size = src0->nb[1]; - const size_t src2_row_size = src2->nb[1]; - - const size_t src0_stride = mmctx->vtcm_src0_stride; - const size_t src2_stride = mmctx->vtcm_src2_stride; - const size_t src1_stride = mmctx->vtcm_src1_stride; - - uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; - uint8_t * restrict vtcm_src2_ptr = mmctx->vtcm_src2 + mmctx->vtcm_src2_size_per_thread * ith; - uint8_t * restrict src1_data = mmctx->vtcm_src1; - - dma_queue * dma_queue = octx->ctx->dma[ith]; - - const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; - const uint32_t n_prefetch = kparams->n_prefetch; - assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); - const uint32_t prefetch_mask = n_prefetch - 1; - - const uint8_t * restrict src0_row = (const uint8_t *) src0->data; - const uint8_t * restrict src2_row = (const uint8_t *) src2->data; - - // Prefill spad with src0, src2 rows - if (src0_start_row < src0_end_row) { for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { const int is0 = (ir0 - src0_start_row); - if (is0 >= (int)n_prefetch) { - break; - } + if (is0 >= (int)n_prefetch) break; dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), src0_stride, src0_row_size, src0_row_size, 2); - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + is0 * src2_stride, src2_row + ir0 * src2_row_size), - src2_stride, src2_row_size, src2_row_size, 2); - } - } - - hvx_mm_run_quant_task(mmctx, ith); - - if (src0_start_row >= src0_end_row) { - return; - } - - // Process rows - for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { - const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; - const uint8_t * ss2 = dma_queue_pop(dma_queue).dst; - - // Process src1 columns in pairs (2×2 tiling) - uint32_t ir1 = 0; - for (; ir1 + 1 < src1_nrows; ir1 += 2) { - const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (ir1+0) * src1_stride); - const uint8_t * restrict src1_col1 = (const uint8_t *) (src1_data + (ir1+1) * src1_stride); - - float * restrict dst_row0_gate = (float *) (dst_gate->data + ((ir1+0) * dst_row_size)); - float * restrict dst_row1_gate = (float *) (dst_gate->data + ((ir1+1) * dst_row_size)); - mmctx->vec_dot_2x2(ne00, &dst_row0_gate[ir0], &dst_row1_gate[ir0], ss0, ss0 + src0_stride, src1_col0, src1_col1); - - float * restrict dst_row0_up = (float *) (dst_up->data + ((ir1+0) * dst_row_size)); - float * restrict dst_row1_up = (float *) (dst_up->data + ((ir1+1) * dst_row_size)); - mmctx->vec_dot_2x2(ne00, &dst_row0_up[ir0], &dst_row1_up[ir0], ss2, ss2 + src2_stride, src1_col0, src1_col1); } - // Handle remaining src1 rows (fallback to 2×1) - for (; ir1 < src1_nrows; ++ir1) { - const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); - - float * restrict dst_row_gate = (float *) (dst_gate->data + (ir1 * dst_row_size)); - mmctx->vec_dot_2x1(ne00, &dst_row_gate[ir0], ss0, ss0 + src0_stride, src1_col); - - float * restrict dst_row_up = (float *) (dst_up->data + (ir1 * dst_row_size)); - mmctx->vec_dot_2x1(ne00, &dst_row_up[ir0], ss2, ss2 + src2_stride, src1_col); - } + for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { + const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0); + uint32_t ir1 = 0; + for (; ir1 + 1 < src1_nrows; ir1 += 2) { + const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (ir1+0) * src1_stride); + const uint8_t * restrict src1_col1 = (const uint8_t *) (src1_data + (ir1+1) * src1_stride); + float * restrict dst_row0 = (float *) (dst->data + ((ir1+0) * dst_row_size)); + float * restrict dst_row1 = (float *) (dst->data + ((ir1+1) * dst_row_size)); + mmctx->vec_dot_2x2(ne00, &dst_row0[ir0], &dst_row1[ir0], ss0, ss0 + src0_stride, src1_col0, src1_col1); + } + for (; ir1 < src1_nrows; ++ir1) { + const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); + float * restrict dst_row = (float *) (dst->data + (ir1 * dst_row_size)); + mmctx->vec_dot_2x1(ne00, &dst_row[ir0], ss0, ss0 + src0_stride, src1_col); + } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0); - // Prefetch next rows - const int pr0 = (ir0 + n_prefetch); - const int is0 = (pr0 - src0_start_row) & prefetch_mask; - if (pr0 < src0_end_row_x2) { - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + pr0 * src0_row_size), - src0_stride, src0_row_size, src0_row_size, 2); - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + is0 * src2_stride, src2_row + pr0 * src2_row_size), - src2_stride, src2_row_size, src2_row_size, 2); + const int pr0 = (ir0 + n_prefetch); + const int is0 = (pr0 - src0_start_row) & prefetch_mask; + if (pr0 < src0_end_row_x2) { + dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + pr0 * src0_row_size), + src0_stride, src0_row_size, src0_row_size, 2); + } } - } - - // Process last row (if any) - if (src0_end_row != src0_end_row_x2) { - uint32_t ir0 = src0_end_row_x2; - const int is0 = (ir0 - src0_start_row) & prefetch_mask; - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), - src0_stride, src0_row_size, src0_row_size, 1); - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + is0 * src2_stride, src2_row + ir0 * src2_row_size), - src2_stride, src2_row_size, src2_row_size, 1); - const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; - const uint8_t * ss2 = dma_queue_pop(dma_queue).dst; - - for (uint32_t ir1 = 0; ir1 < src1_nrows; ++ir1) { - const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); - - float * restrict dst_row_gate = (float *) (dst_gate->data + (ir1 * dst_row_size)); - mmctx->vec_dot_1x1(ne00, &dst_row_gate[ir0], ss0, src1_col); - - float * restrict dst_row_up = (float *) (dst_up->data + (ir1 * dst_row_size)); - mmctx->vec_dot_1x1(ne00, &dst_row_up[ir0], ss2, src1_col); + if (src0_end_row != src0_end_row_x2) { + uint32_t ir0 = src0_end_row_x2; + const int is0 = (ir0 - src0_start_row) & prefetch_mask; + dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), + src0_stride, src0_row_size, src0_row_size, 1); + const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0); + for (uint32_t ir1 = 0; ir1 < src1_nrows; ++ir1) { + const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); + float * restrict dst_row = (float *) (dst->data + (ir1 * dst_row_size)); + mmctx->vec_dot_1x1(ne00, &dst_row[ir0], ss0, src1_col); + } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0); } } } @@ -3485,7 +3074,7 @@ static int hvx_mm_matmul_id( struct htp_mm_hvx_vtcm_layout L; htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads, - 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, true, false, false); + 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, true, false); size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes; @@ -3610,6 +3199,7 @@ int op_matmul_id(struct htp_ops_context * octx) { struct htp_mm_context mmctx_struct = {0}; struct htp_mm_context * mmctx = &mmctx_struct; mmctx->octx = octx; + mmctx->act = src1; const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; @@ -3690,18 +3280,15 @@ int op_matmul_id(struct htp_ops_context * octx) { return s; } - -int op_matmul_qkv(struct htp_ops_context * octx) { +int op_matmul_nx(struct htp_ops_context * octx) { struct htp_thread_trace * tr = &octx->ctx->trace[0]; htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); - const struct htp_tensor * restrict src0 = octx->src[0]; // Wk - const struct htp_tensor * restrict src1 = octx->src[1]; // x - const struct htp_tensor * restrict src2 = octx->src[2]; // Wv - const struct htp_tensor * restrict src3 = octx->src[3]; // Wq - const struct htp_tensor * restrict dst_k = octx->dsts[0]; - const struct htp_tensor * restrict dst_v = octx->dsts[1]; - const struct htp_tensor * restrict dst_q = octx->dsts[2]; + const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; + const uint32_t n_weights = kparams->n_weights; + + const struct htp_tensor * restrict src0 = octx->src[0]; // first weight + const struct htp_tensor * restrict act = octx->src[n_weights]; // activation x bool is_repacked = (src0->type == HTP_TYPE_Q4_0 || src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q8_0 || src0->type == HTP_TYPE_IQ4_NL || @@ -3710,19 +3297,9 @@ int op_matmul_qkv(struct htp_ops_context * octx) { struct htp_mm_context mmctx_struct = {0}; struct htp_mm_context * mmctx = &mmctx_struct; mmctx->octx = octx; + mmctx->act = act; - const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; - - const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3]; - const uint32_t src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3]; - - // Compute src0_nrows_per_thread - mmctx->src0_nrows_per_thread = (src0_nrows + octx->n_threads - 1) / octx->n_threads; - if (is_repacked) { - mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32); - } else { - mmctx->src0_nrows_per_thread += (mmctx->src0_nrows_per_thread & 1); // round up to even - } + const uint32_t src1_nrows = act->ne[1] * act->ne[2] * act->ne[3]; const size_t src0_row_size = src0->nb[1]; const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128); @@ -3732,7 +3309,7 @@ int op_matmul_qkv(struct htp_ops_context * octx) { } const uint32_t qk = QK_Q8_0_TILED; - const uint32_t nb = (src1->ne[0] + qk - 1) / qk; + const uint32_t nb = (act->ne[0] + qk - 1) / qk; const uint32_t total_nb = src1_nrows * nb; worker_callback_t quant_task_func; @@ -3758,185 +3335,39 @@ int op_matmul_qkv(struct htp_ops_context * octx) { size_t src1_row_size; if (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { - src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(src1->ne[0]) : htp_mm_q8_0_flat_row_size(src1->ne[0]); + src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(act->ne[0]) : htp_mm_q8_0_flat_row_size(act->ne[0]); } else { - src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(src1->ne[0]) : htp_mm_q8_0_tiled_row_size(src1->ne[0]); + src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(act->ne[0]) : htp_mm_q8_0_tiled_row_size(act->ne[0]); } struct htp_mm_hvx_vtcm_layout L; - htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, src1->ne[0], src1_nrows, octx->n_threads, - 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, false, true, false); + htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, act->ne[0], src1_nrows, octx->n_threads, + 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, false, true); size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes; if (octx->ctx->vtcm_size < vtcm_size) { - FARF(ERROR, "matmul-qkv: current VTCM reservation %zu is too small, needed %zu\n", + FARF(ERROR, "matmul-nx: current VTCM reservation %zu is too small, needed %zu\n", octx->ctx->vtcm_size, vtcm_size); return HTP_STATUS_VTCM_TOO_SMALL; } uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base; - mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1); mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0); - mmctx->vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src2); - mmctx->vtcm_src3 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src3); - mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); - - octx->src1_spad.src = NULL; - octx->src0_spad.src = NULL; - octx->src2_spad.src = NULL; - octx->src3_spad.src = NULL; - octx->dst_spad.src = NULL; - - mmctx->vtcm_src0_stride = is_repacked ? 0 : src0_row_size_padded; - mmctx->vtcm_src2_stride = is_repacked ? 0 : src0_row_size_padded; - mmctx->vtcm_src3_stride = is_repacked ? 0 : src0_row_size_padded; - mmctx->vtcm_src1_stride = src1_row_size; - - mmctx->vtcm_src0_size_per_thread = L.src0_bytes / octx->n_threads; - mmctx->vtcm_src1_size_per_thread = L.src1_bytes; - mmctx->vtcm_src2_size_per_thread = L.src2_bytes / octx->n_threads; - mmctx->vtcm_src3_size_per_thread = L.src3_bytes / octx->n_threads; - mmctx->vtcm_dst_size_per_thread = L.dst_bytes / octx->n_threads; - - mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks; - mmctx->quant_task_func = quant_task_func; - mmctx->n_quant_tasks = n_quant_tasks; - atomic_init(&mmctx->quant_barrier, n_quant_tasks); - - // Run fused matmul - const uint32_t n_matmul_jobs = octx->n_threads; - worker_callback_t matmul_job_func; - if (is_repacked) { - if (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { - switch (src0->type) { - case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_qkv_2d_repacked_q4_0_flat; break; - case HTP_TYPE_Q4_1: matmul_job_func = hvx_mm_qkv_2d_repacked_q4_1_flat; break; - case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_qkv_2d_repacked_q8_0_flat; break; - case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_qkv_2d_repacked_iq4nl_flat; break; - case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_qkv_2d_repacked_mxfp4_flat; break; - default: return HTP_STATUS_NO_SUPPORT; - } - } else { - switch (src0->type) { - case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_qkv_2d_repacked_q4_0; break; - case HTP_TYPE_Q4_1: matmul_job_func = hvx_mm_qkv_2d_repacked_q4_1; break; - case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_qkv_2d_repacked_q8_0; break; - case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_qkv_2d_repacked_iq4nl; break; - case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_qkv_2d_repacked_mxfp4; break; - default: return HTP_STATUS_NO_SUPPORT; - } - } - } else { - matmul_job_func = hvx_mm_qkv_2d; - } - - htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); - - worker_pool_run_func(octx->ctx->worker_pool, matmul_job_func, mmctx, n_matmul_jobs); - - return HTP_STATUS_OK; -} - -int op_matmul_ffn(struct htp_ops_context * octx) { - struct htp_thread_trace * tr = &octx->ctx->trace[0]; - htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); - - const struct htp_tensor * restrict src0 = octx->src[0]; // Wgate - const struct htp_tensor * restrict src1 = octx->src[1]; // y - const struct htp_tensor * restrict src2 = octx->src[2]; // Wup - const struct htp_tensor * restrict dst_gate = octx->dsts[0]; - const struct htp_tensor * restrict dst_up = octx->dsts[1]; - - bool is_repacked = (src0->type == HTP_TYPE_Q4_0 || src0->type == HTP_TYPE_Q4_1 || - src0->type == HTP_TYPE_Q8_0 || src0->type == HTP_TYPE_IQ4_NL || - src0->type == HTP_TYPE_MXFP4); - - struct htp_mm_context mmctx_struct = {0}; - struct htp_mm_context * mmctx = &mmctx_struct; - mmctx->octx = octx; - - const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; - - const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3]; - const uint32_t src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3]; - - // Compute src0_nrows_per_thread - mmctx->src0_nrows_per_thread = (src0_nrows + octx->n_threads - 1) / octx->n_threads; - if (is_repacked) { - mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32); - } else { - mmctx->src0_nrows_per_thread += (mmctx->src0_nrows_per_thread & 1); // round up to even - } - - const size_t src0_row_size = src0->nb[1]; - const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128); - - if (hvx_mm_init_vec_dot(mmctx, src0->type) != 0) { - return HTP_STATUS_NO_SUPPORT; - } - - const uint32_t qk = QK_Q8_0_TILED; - const uint32_t nb = (src1->ne[0] + qk - 1) / qk; - const uint32_t total_nb = src1_nrows * nb; - - worker_callback_t quant_task_func; - uint32_t n_quant_tasks = 1; - if (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { - n_quant_tasks = MIN(src1_nrows, octx->n_threads); - quant_task_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_flat : quantize_f32_q8_0_flat; - } else if (src1_nrows < octx->n_threads) { - n_quant_tasks = MIN(total_nb, octx->n_threads); - quant_task_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block; - for (uint32_t ith = 0; ith < n_quant_tasks; ++ith) { - uint32_t ib_first = (total_nb * (ith + 0)) / n_quant_tasks; - uint32_t ib_last = (total_nb * (ith + 1)) / n_quant_tasks; - mmctx->quant_ib_first[ith] = ib_first; - mmctx->quant_ib_last[ith] = ib_last; - mmctx->quant_r[ith] = ib_first / nb; - mmctx->quant_c[ith] = ib_first % nb; - } - } else { - n_quant_tasks = MIN(src1_nrows, octx->n_threads); - quant_task_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled : quantize_f32_q8_0_tiled; - } - - size_t src1_row_size; - if (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { - src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(src1->ne[0]) : htp_mm_q8_0_flat_row_size(src1->ne[0]); - } else { - src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(src1->ne[0]) : htp_mm_q8_0_tiled_row_size(src1->ne[0]); - } - - struct htp_mm_hvx_vtcm_layout L; - htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, src1->ne[0], src1_nrows, octx->n_threads, - 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, false, false, true); - - size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes; - - if (octx->ctx->vtcm_size < vtcm_size) { - FARF(ERROR, "matmul-ffn: current VTCM reservation %zu is too small, needed %zu\n", octx->ctx->vtcm_size, vtcm_size); - return HTP_STATUS_VTCM_TOO_SMALL; - } - - uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base; mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1); - mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0); - mmctx->vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src2); mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); - octx->src1_spad.src = NULL; octx->src0_spad.src = NULL; + octx->src1_spad.src = NULL; octx->src2_spad.src = NULL; + octx->src3_spad.src = NULL; octx->dst_spad.src = NULL; mmctx->vtcm_src0_stride = is_repacked ? 0 : src0_row_size_padded; - mmctx->vtcm_src2_stride = is_repacked ? 0 : src0_row_size_padded; mmctx->vtcm_src1_stride = src1_row_size; mmctx->vtcm_src0_size_per_thread = L.src0_bytes / octx->n_threads; mmctx->vtcm_src1_size_per_thread = L.src1_bytes; - mmctx->vtcm_src2_size_per_thread = L.src2_bytes / octx->n_threads; mmctx->vtcm_dst_size_per_thread = L.dst_bytes / octx->n_threads; mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks; @@ -3950,25 +3381,25 @@ int op_matmul_ffn(struct htp_ops_context * octx) { if (is_repacked) { if (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { switch (src0->type) { - case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_ffn_2d_repacked_q4_0_flat; break; - case HTP_TYPE_Q4_1: matmul_job_func = hvx_mm_ffn_2d_repacked_q4_1_flat; break; - case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_ffn_2d_repacked_q8_0_flat; break; - case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_ffn_2d_repacked_iq4nl_flat; break; - case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_ffn_2d_repacked_mxfp4_flat; break; + case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_nx_2d_repacked_q4_0_flat; break; + case HTP_TYPE_Q4_1: matmul_job_func = hvx_mm_nx_2d_repacked_q4_1_flat; break; + case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_nx_2d_repacked_q8_0_flat; break; + case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_nx_2d_repacked_iq4nl_flat; break; + case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_nx_2d_repacked_mxfp4_flat; break; default: return HTP_STATUS_NO_SUPPORT; } } else { switch (src0->type) { - case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_ffn_2d_repacked_q4_0; break; - case HTP_TYPE_Q4_1: matmul_job_func = hvx_mm_ffn_2d_repacked_q4_1; break; - case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_ffn_2d_repacked_q8_0; break; - case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_ffn_2d_repacked_iq4nl; break; - case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_ffn_2d_repacked_mxfp4; break; + case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_nx_2d_repacked_q4_0; break; + case HTP_TYPE_Q4_1: matmul_job_func = hvx_mm_nx_2d_repacked_q4_1; break; + case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_nx_2d_repacked_q8_0; break; + case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_nx_2d_repacked_iq4nl; break; + case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_nx_2d_repacked_mxfp4; break; default: return HTP_STATUS_NO_SUPPORT; } } } else { - matmul_job_func = hvx_mm_ffn_2d; + matmul_job_func = hvx_mm_nx_2d; } htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.h b/ggml/src/ggml-hexagon/htp/matmul-ops.h index 6c393664c6e..dbc8e359093 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.h +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.h @@ -88,6 +88,7 @@ struct htp_mm_kernel_params { int32_t vtcm_src2_size; // src2 scratchpad size in VTCM (fused only) int32_t vtcm_src3_size; // src3 scratchpad size in VTCM (fused only) int32_t vtcm_dst_size; // dst scratchpad size in VTCM + int32_t n_weights; // Number of weights for fused NX // Precomputed division values struct fastdiv_values div_ne12_ne1; @@ -463,8 +464,7 @@ static inline void htp_mm_hvx_vtcm_layout_build( size_t src2_row_size, uint32_t n_prefetch, bool is_matmul_id, - bool is_fused_qkv, - bool is_fused_ffn + bool is_fused_nx ) { size_t src0_sz = 0; size_t src1_sz = 0; @@ -476,44 +476,33 @@ static inline void htp_mm_hvx_vtcm_layout_build( wtype == HTP_TYPE_Q8_0 || wtype == HTP_TYPE_IQ4_NL || wtype == HTP_TYPE_MXFP4); - if (is_fused_qkv || is_fused_ffn) { + if (is_fused_nx) { const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128); const size_t quant_scratch_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * n_threads; - size_t src0_sz_per_thread = 0; - size_t src2_sz_per_thread = 0; - size_t src3_sz_per_thread = 0; + size_t weight_sz_per_thread = 0; if (is_repack) { uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype); uint32_t n_k_tiles = hex_round_up(ne10, 32) / 32; uint32_t tile_row_size = n_k_tiles * aligned_tile_size; - src0_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128); - src2_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128); - if (is_fused_qkv) { - src3_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128); - } + weight_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128); } else { - src0_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128); - src2_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128); - if (is_fused_qkv) { - src3_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128); - } + weight_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128); } - size_t flat_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); - size_t tiled_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); + size_t flat_act_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); + size_t tiled_act_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); - if (kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { - src1_sz = hex_round_up(flat_src1_row_size * src1_nrows, 128); - } else { - src1_sz = hex_round_up(tiled_src1_row_size * src1_nrows, 128); - } + size_t act_sz = (kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) + ? hex_round_up(flat_act_row_size * src1_nrows, 128) + : hex_round_up(tiled_act_row_size * src1_nrows, 128); - src0_sz = src0_sz_per_thread * n_threads; - src2_sz = src2_sz_per_thread * n_threads; - src3_sz = src3_sz_per_thread * n_threads; + src0_sz = weight_sz_per_thread * n_threads; // shared single-weight prefetch buffer + src1_sz = act_sz; // quantized activation buffer + src2_sz = 0; + src3_sz = 0; dst_sz = quant_scratch_size; } else if (is_matmul_id) { const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128); @@ -616,8 +605,8 @@ static inline void htp_mm_hvx_vtcm_layout_build( } size_t off = 0; - VTCM_LAYOUT_ALLOC(off, off_src1, src1_sz); VTCM_LAYOUT_ALLOC(off, off_src0, src0_sz); + VTCM_LAYOUT_ALLOC(off, off_src1, src1_sz); VTCM_LAYOUT_ALLOC(off, off_src2, src2_sz); VTCM_LAYOUT_ALLOC(off, off_src3, src3_sz); VTCM_LAYOUT_ALLOC(off, off_dst, dst_sz); diff --git a/ggml/src/ggml-hexagon/htp/set-rows-ops.c b/ggml/src/ggml-hexagon/htp/set-rows-ops.c index 58c54967db0..fa14bf0ef6b 100644 --- a/ggml/src/ggml-hexagon/htp/set-rows-ops.c +++ b/ggml/src/ggml-hexagon/htp/set-rows-ops.c @@ -8,14 +8,20 @@ #include #include -#include "hex-dma.h" +#include "dma-queue.h" +#include "work-queue.h" #include "hvx-utils.h" +#include "hex-utils.h" +#include "hvx-copy.h" +#include "hvx-quant.h" #define GGML_COMMON_DECL_C #include "ggml-common.h" + #include "htp-ctx.h" #include "htp-ops.h" -#include "htp-ops.h" +#include "htp-tensor.h" +#include "htp/set-rows-ops.h" #define set_rows_preamble \ const uint32_t ne00 = octx->src[0]->ne[0]; \ @@ -47,116 +53,142 @@ \ const uint32_t nr = ne01; -struct htp_set_rows_context { +struct set_rows_context { struct htp_ops_context * octx; - struct fastdiv_values div_ne12; - struct fastdiv_values div_ne11; - uint32_t src0_nrows_per_thread; + const struct htp_set_rows_kernel_params * kparams; + struct htp_set_rows_vtcm_layout vtcm_layout; + uint8_t * vtcm_base; }; -static void set_rows_thread_f32_f32(unsigned int nth, unsigned int ith, void *data) { - struct htp_set_rows_context * srctx = (struct htp_set_rows_context *)data; - struct htp_ops_context * octx = srctx->octx; - - set_rows_preamble; - - uint64_t qt = HAP_perf_get_qtimer_count(); - - // parallelize by rows of src0 - const uint32_t dr = srctx->src0_nrows_per_thread; - const uint32_t ir0 = dr * ith; - if (ir0 >= nr) { - return; - } - const uint32_t ir1 = (ir0 + dr < nr) ? (ir0 + dr) : nr; - - const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32); - - for (uint32_t i03 = 0; i03 < ne03; ++i03) { - for (uint32_t i02 = 0; i02 < ne02; ++i02) { - for (uint32_t i = ir0; i < ir1; ++i) { - const uint32_t i12 = fastmodulo(i03, ne12, &srctx->div_ne12); - const uint32_t i11 = fastmodulo(i02, ne11, &srctx->div_ne11); - const uint32_t i10 = i; - - const uintptr_t src1_addr = octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12; - - uint32_t i1 = is_i32 ? *(int32_t *)src1_addr : *(int64_t *)src1_addr; - if (i1 >= ne1) { - // ignore invalid indices - continue; - } - - const uintptr_t src0_ptr = octx->src[0]->data + i*nb01 + i02*nb02 + i03*nb03; - const uintptr_t dst_ptr = octx->dst->data + i1*nb1 + i02*nb2 + i03*nb3; - - // copy row - hvx_copy_f32_uu((uint8_t *)dst_ptr, (const uint8_t *)src0_ptr, ne00); - } - } - } - - qt = HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - qt); - FARF(HIGH, "set-rows-f32-f32 %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u usec %u\n", ith, nth, - ne00, ne01, ne02, ne03, ir0, ir1, ne10, ne11, ne12, ne13, ne0, ne1, ne2, ne3, (unsigned) qt); +#define SET_ROWS_THREAD_DMA_FN(TYPE_NAME, IDX_TYPE, COMPUTE_EXPR) \ +static void set_rows_thread_dma_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsigned int ith, void *data) { \ + struct set_rows_context * srctx = (struct set_rows_context *)data; \ + struct htp_ops_context * octx = srctx->octx; \ + const struct htp_set_rows_kernel_params * kparams = srctx->kparams; \ + set_rows_preamble; \ + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ + const uint32_t dr = kparams->tasks_per_thread; \ + const uint32_t ir0 = dr * ith; \ + if (ir0 >= kparams->total_tasks) { \ + return; \ + } \ + const uint32_t ir1 = MIN(ir0 + dr, kparams->total_tasks); \ + dma_queue * dma_queue = octx->ctx->dma[ith]; \ + const struct htp_set_rows_vtcm_layout * vtcm_layout = &srctx->vtcm_layout; \ + uint8_t * vtcm_src0 = srctx->vtcm_base + vtcm_layout->off_src0 + ith * vtcm_layout->src0_bytes_per_thread; \ + uint8_t * vtcm_dst = srctx->vtcm_base + vtcm_layout->off_dst + ith * vtcm_layout->dst_bytes_per_thread; \ + const uint32_t src0_row_size = ne00 * sizeof(float); \ + const uint32_t dst_row_size = htp_tensor_get_row_size(octx->dst->type, ne00); \ + const uint32_t nrows_per_thread = ir1 - ir0; \ + const uint32_t total_steps = ne03 * ne02 * nrows_per_thread; \ + uint32_t pi_step = 0; \ + uint32_t pi02 = 0; \ + uint32_t pi03 = 0; \ + for (uint32_t step = 0, spad_idx = 0; step < total_steps && spad_idx < 2; ++step, spad_idx++) { \ + uint32_t i = ir0 + pi_step; \ + const uintptr_t src0_ptr = octx->src[0]->data + i*nb01 + pi02*nb02 + pi03*nb03; \ + dma_queue_push(dma_queue, \ + dma_make_ptr((void *)octx->dst->data, \ + vtcm_dst + spad_idx * vtcm_layout->dst_spad_half_size), \ + dst_row_size, vtcm_layout->dst_spad_half_size, dst_row_size, 0); \ + dma_queue_push(dma_queue, \ + dma_make_ptr((void *)(vtcm_src0 + spad_idx * vtcm_layout->src0_spad_half_size), \ + (const void *)src0_ptr), \ + vtcm_layout->src0_spad_half_size, src0_row_size, src0_row_size, 1); \ + pi_step++; \ + if (pi_step == nrows_per_thread) { \ + pi_step = 0; \ + pi02++; \ + if (pi02 == ne02) { \ + pi02 = 0; \ + pi03++; \ + } \ + } \ + } \ + uint32_t ci_step = 0; \ + uint32_t ci02 = 0; \ + uint32_t ci03 = 0; \ + uint32_t ci11_base = 0; \ + uint32_t ci12_base = 0; \ + for (uint32_t step = 0; step < total_steps; ++step) { \ + void * dst_spad = (void *) dma_queue_pop(dma_queue).src; \ + void * src_spad = (void *) dma_queue_pop(dma_queue).dst; \ + uint32_t i = ir0 + ci_step; \ + const uintptr_t src1_addr = octx->src[1]->data + i*nb10 + ci11_base*nb11 + ci12_base*nb12; \ + const IDX_TYPE i1 = *(const IDX_TYPE *)src1_addr; \ + const bool valid_i1 = ((uint64_t)i1 < (uint64_t)ne1); \ + const uint32_t target_i1 = (uint32_t)i1; \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, step); \ + if (valid_i1) { \ + COMPUTE_EXPR; \ + } \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, step); \ + if (valid_i1) { \ + const uintptr_t dst_ptr = octx->dst->data + target_i1*nb1 + ci02*nb2 + ci03*nb3; \ + dma_queue_push(dma_queue, \ + dma_make_ptr((void *)dst_ptr, (const void *)dst_spad), \ + dst_row_size, vtcm_layout->dst_spad_half_size, dst_row_size, 1); \ + } else { \ + dma_queue_push(dma_queue, \ + dma_make_ptr((void *)octx->dst->data, (const void *)dst_spad), \ + dst_row_size, vtcm_layout->dst_spad_half_size, dst_row_size, 0); \ + } \ + const uint32_t next_step = step + 2; \ + if (next_step < total_steps) { \ + uint32_t ni = ir0 + pi_step; \ + const uintptr_t psrc0_ptr = octx->src[0]->data + ni*nb01 + pi02*nb02 + pi03*nb03; \ + dma_queue_push(dma_queue, \ + dma_make_ptr((void *)src_spad, (const void *)psrc0_ptr), \ + vtcm_layout->src0_spad_half_size, src0_row_size, src0_row_size, 1); \ + pi_step++; \ + if (pi_step == nrows_per_thread) { \ + pi_step = 0; \ + pi02++; \ + if (pi02 == ne02) { \ + pi02 = 0; \ + pi03++; \ + } \ + } \ + } \ + ci_step++; \ + if (ci_step == nrows_per_thread) { \ + ci_step = 0; \ + ci02++; \ + ci11_base++; \ + if (ci11_base == ne11) { \ + ci11_base = 0; \ + } \ + if (ci02 == ne02) { \ + ci02 = 0; \ + ci03++; \ + ci12_base++; \ + if (ci12_base == ne12) { \ + ci12_base = 0; \ + } \ + } \ + } \ + } \ + dma_queue_flush(dma_queue); \ } -static void set_rows_thread_f16_f32(unsigned int nth, unsigned int ith, void *data) { - struct htp_set_rows_context * srctx = (struct htp_set_rows_context *)data; - struct htp_ops_context * octx = srctx->octx; - - set_rows_preamble; +SET_ROWS_THREAD_DMA_FN(f32, int32_t, { hvx_copy_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, ne00); }) +SET_ROWS_THREAD_DMA_FN(f32, int64_t, { hvx_copy_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, ne00); }) - uint64_t qt = HAP_perf_get_qtimer_count(); +SET_ROWS_THREAD_DMA_FN(f16, int32_t, { hvx_copy_f16_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, ne00); }) +SET_ROWS_THREAD_DMA_FN(f16, int64_t, { hvx_copy_f16_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, ne00); }) - // parallelize by rows of src0 - const uint32_t dr = srctx->src0_nrows_per_thread; - const uint32_t ir0 = dr * ith; - if (ir0 >= nr) { - return; - } - const uint32_t ir1 = (ir0 + dr < nr) ? (ir0 + dr) : nr; - - const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32); - - for (uint32_t i03 = 0; i03 < ne03; ++i03) { - for (uint32_t i02 = 0; i02 < ne02; ++i02) { - for (uint32_t i = ir0; i < ir1; ++i) { - const uint32_t i12 = fastmodulo(i03, ne12, &srctx->div_ne12); - const uint32_t i11 = fastmodulo(i02, ne11, &srctx->div_ne11); - const uint32_t i10 = i; - - const uintptr_t src1_addr = octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12; - - uint32_t i1 = is_i32 ? *(int32_t *)src1_addr : *(int64_t *)src1_addr; - if (i1 >= ne1) { - // ignore invalid indices - continue; - } - - const uint8_t* src0_ptr = (const uint8_t *) octx->src[0]->data + i*nb01 + i02*nb02 + i03*nb03; - uint8_t* dst_ptr = (uint8_t *) octx->dst->data + i1*nb1 + i02*nb2 + i03*nb3; - - hvx_copy_f16_f32_uu(dst_ptr, src0_ptr, ne00); - } - } - } - - qt = HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - qt); - FARF(HIGH, "set-rows-f16-f32 %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u usec %u\n", ith, nth, - ne00, ne01, ne02, ne03, ir0, ir1, ne10, ne11, ne12, ne13, ne0, ne1, ne2, ne3, (unsigned) qt); -} +SET_ROWS_THREAD_DMA_FN(q8_0, int32_t, { hvx_quantize_row_q8_0_f32(dst_spad, (const float *)src_spad, ne00); }) +SET_ROWS_THREAD_DMA_FN(q8_0, int64_t, { hvx_quantize_row_q8_0_f32(dst_spad, (const float *)src_spad, ne00); }) int op_set_rows(struct htp_ops_context * octx) { + const struct htp_set_rows_kernel_params * kparams = (const struct htp_set_rows_kernel_params *)octx->kernel_params; set_rows_preamble; - const uint32_t n_threads = MIN(nr, octx->n_threads); - if (octx->src[0]->type != HTP_TYPE_F32) { return HTP_STATUS_NO_SUPPORT; } - if (octx->dst->type != HTP_TYPE_F32 && octx->dst->type != HTP_TYPE_F16) { + if (octx->dst->type != HTP_TYPE_F32 && octx->dst->type != HTP_TYPE_F16 && octx->dst->type != HTP_TYPE_Q8_0) { return HTP_STATUS_NO_SUPPORT; } @@ -164,27 +196,27 @@ int op_set_rows(struct htp_ops_context * octx) { return HTP_STATUS_NO_SUPPORT; } - if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) { - return HTP_STATUS_OK; - } + // l2fetch the src1 (indices) tensor in the main thread + hex_l2fetch_block((const void *)octx->src[1]->data, octx->src[1]->ne[3] * octx->src[1]->nb[3]); - struct htp_set_rows_context srctx; + struct set_rows_context srctx; srctx.octx = octx; - srctx.div_ne12 = init_fastdiv_values(ne12); - srctx.div_ne11 = init_fastdiv_values(ne11); - - srctx.src0_nrows_per_thread = (nr + n_threads - 1) / n_threads; - - switch(octx->dst->type) { - case HTP_TYPE_F32: - worker_pool_run_func(octx->ctx->worker_pool, set_rows_thread_f32_f32, &srctx, n_threads); - break; - case HTP_TYPE_F16: - worker_pool_run_func(octx->ctx->worker_pool, set_rows_thread_f16_f32, &srctx, n_threads); - break; - default: - return HTP_STATUS_NO_SUPPORT; + srctx.kparams = kparams; + + htp_set_rows_vtcm_layout_build(&srctx.vtcm_layout, octx->dst->type, ne00, kparams->n_threads); + srctx.vtcm_base = (uint8_t *)octx->ctx->vtcm_base; + + work_queue_func_t q_func = NULL; + const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32); + + switch (octx->dst->type) { + case HTP_TYPE_F32: q_func = is_i32 ? set_rows_thread_dma_f32_int32_t : set_rows_thread_dma_f32_int64_t; break; + case HTP_TYPE_F16: q_func = is_i32 ? set_rows_thread_dma_f16_int32_t : set_rows_thread_dma_f16_int64_t; break; + case HTP_TYPE_Q8_0: q_func = is_i32 ? set_rows_thread_dma_q8_0_int32_t : set_rows_thread_dma_q8_0_int64_t; break; + default: return HTP_STATUS_NO_SUPPORT; } + work_queue_run(octx->ctx->work_queue, q_func, &srctx, kparams->n_threads); + return HTP_STATUS_OK; } diff --git a/ggml/src/ggml-hexagon/htp/set-rows-ops.h b/ggml/src/ggml-hexagon/htp/set-rows-ops.h new file mode 100644 index 00000000000..5e98d2cb55c --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/set-rows-ops.h @@ -0,0 +1,74 @@ +#ifndef HTP_SET_ROWS_OPS_H +#define HTP_SET_ROWS_OPS_H + +#include "hex-fastdiv.h" + +struct htp_set_rows_kernel_params { + int32_t n_threads; + int32_t total_tasks; + int32_t tasks_per_thread; + int32_t vtcm_size; + + // Fastdiv helpers + struct fastdiv_values div_ne11; + struct fastdiv_values div_ne12; + struct fastdiv_values div_tasks_per_thread; + struct fastdiv_values div_ne02; +}; + +struct htp_set_rows_vtcm_layout { + size_t total_bytes; + size_t off_src0; + size_t off_dst; + + size_t src0_bytes_per_thread; + size_t dst_bytes_per_thread; + + size_t src0_spad_half_size; + size_t dst_spad_half_size; +}; + +static inline void htp_set_rows_vtcm_layout_build( + struct htp_set_rows_vtcm_layout * vtcm_layout, + int dst_type, + uint32_t ne00, + uint32_t n_threads) { + + size_t src0_row_size = ne00 * 4; + size_t dst_row_size = 0; + switch (dst_type) { + case 0: // HTP_TYPE_F32 + dst_row_size = ne00 * 4; + break; + case 1: // HTP_TYPE_F16 + dst_row_size = ne00 * 2; + break; + case 8: // HTP_TYPE_Q8_0 + dst_row_size = (ne00 / 32) * 34; + break; + default: + dst_row_size = 0; + break; + } + + size_t src0_row_size_aligned = (src0_row_size + 255) & ~255; + size_t dst_row_size_aligned = (dst_row_size + 255) & ~255; + + vtcm_layout->src0_spad_half_size = src0_row_size_aligned; + vtcm_layout->dst_spad_half_size = dst_row_size_aligned; + + vtcm_layout->src0_bytes_per_thread = src0_row_size_aligned * 2; + vtcm_layout->dst_bytes_per_thread = dst_row_size_aligned * 2; + + vtcm_layout->off_src0 = 0; + vtcm_layout->off_dst = vtcm_layout->off_src0 + vtcm_layout->src0_bytes_per_thread * n_threads; + vtcm_layout->total_bytes = vtcm_layout->off_dst + vtcm_layout->dst_bytes_per_thread * n_threads; +} + +#if defined(__cplusplus) +static_assert(sizeof(struct htp_set_rows_kernel_params) <= 128, "htp_set_rows_kernel_params is too large for kernel_params blob"); +#else +_Static_assert(sizeof(struct htp_set_rows_kernel_params) <= 128, "htp_set_rows_kernel_params is too large for kernel_params blob"); +#endif + +#endif // HTP_SET_ROWS_OPS_H diff --git a/ggml/src/ggml-hexagon/htp/unary-ops.c b/ggml/src/ggml-hexagon/htp/unary-ops.c index b21415a67d6..1a632bf5631 100644 --- a/ggml/src/ggml-hexagon/htp/unary-ops.c +++ b/ggml/src/ggml-hexagon/htp/unary-ops.c @@ -443,6 +443,34 @@ static void tanh_f32(const float * restrict src, } } +static void abs_f32(const float * restrict src, + float * restrict dst, + const uint32_t num_rows, + const struct htp_unary_context * uctx) { + htp_unary_op_preamble; + + for (uint32_t ir = 0; ir < num_rows; ir++) { + const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned); + uint8_t * restrict dst_local = (uint8_t *)dst + (ir * dst_row_size_aligned); + + hvx_abs_f32_aa(dst_local, src_local, ne0); + } +} + +static void log_f32(const float * restrict src, + float * restrict dst, + const uint32_t num_rows, + const struct htp_unary_context * uctx) { + htp_unary_op_preamble; + + for (uint32_t ir = 0; ir < num_rows; ir++) { + const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned); + uint8_t * restrict dst_local = (uint8_t *)dst + (ir * dst_row_size_aligned); + + hvx_log_f32_aa(dst_local, src_local, ne0); + } +} + #define DEFINE_UNARY_TASK(NAME, IS_RMS_NORM_MUL, IS_TRI, CORE_EXPR) \ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * data) { \ const struct htp_unary_context * uctx = (const struct htp_unary_context *) data; \ @@ -478,6 +506,9 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat const uint32_t nb11 = src1 ? src1->nb[1] : 0; \ const uint32_t nb12 = src1 ? src1->nb[2] : 0; \ const uint32_t nb13 = src1 ? src1->nb[3] : 0; \ + const uint32_t nb11_bc = (src1 && src1->ne[1] > 1) ? nb11 : 0; \ + const uint32_t nb12_bc = (src1 && src1->ne[2] > 1) ? nb12 : 0; \ + const uint32_t nb13_bc = (src1 && src1->ne[3] > 1) ? nb13 : 0; \ const bool src1_contig = src1 ? ((nb12 == (size_t)ne01 * nb11) && (nb13 == (size_t)ne02 * nb12)) : false; \ \ uint8_t * src0_vtcm_data = uctx->vtcm_src0 + (ith * uctx->vtcm_src0_size_per_thread); \ @@ -497,8 +528,12 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat const struct fastdiv_values * div_ne02 = &uctx->kparams->div_ne02; \ const struct fastdiv_values * div_ne012 = &uctx->kparams->div_ne012; \ \ - const uint32_t src0_max_block = src0_contig ? uctx->block : MIN((uint32_t)uctx->block, ne01); \ - const uint32_t dst_max_block = dst_contig ? uctx->block : MIN((uint32_t)uctx->block, ne1); \ + const bool src1_needs_row_clip = (IS_RMS_NORM_MUL) && !uctx->broadcast_weight && !src1_contig; \ + const bool block_src0_contig = src0_contig && !src1_needs_row_clip; \ + const bool block_dst_contig = dst_contig && !src1_needs_row_clip; \ + \ + const uint32_t src0_max_block = block_src0_contig ? uctx->block : MIN((uint32_t)uctx->block, ne01); \ + const uint32_t dst_max_block = block_dst_contig ? uctx->block : MIN((uint32_t)uctx->block, ne1); \ const uint32_t BLOCK = MIN(src0_max_block, dst_max_block); \ if (BLOCK == 0) { \ FARF(ERROR, "unary-f32 : current VTCM reservation %zu is too small, needed at least %zu\n", \ @@ -515,8 +550,8 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat } \ \ for (uint32_t ir = src0_start_row, vtcm_idx = 0; ir < src0_end_row && vtcm_idx < 2; vtcm_idx++) { \ - const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, \ - div_ne01); \ + const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, block_src0_contig, block_dst_contig, \ + ne01, div_ne01); \ \ dma_queue_push(dma_queue, \ dma_make_ptr(data_dst, dst_vtcm_data + (vtcm_idx * dst_vtcm_half_size)), \ @@ -530,7 +565,7 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat \ if ((IS_RMS_NORM_MUL) && !uctx->broadcast_weight) { \ const size_t src1_off = src1_contig ? (ir * nb11) : \ - unary_row_offset(ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11, nb12, nb13); \ + unary_row_offset(ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11_bc, nb12_bc, nb13_bc); \ dma_queue_push(dma_queue, \ dma_make_ptr(src1_vtcm_data + (vtcm_idx * src1_vtcm_half_size), data_src1 + src1_off), \ uctx->src1_row_size_aligned, nb11, uctx->src1_data_row_size, block_size); \ @@ -540,8 +575,8 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat } \ \ for (uint32_t ir = src0_start_row; ir < src0_end_row; ) { \ - const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, \ - div_ne01); \ + const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, block_src0_contig, block_dst_contig, \ + ne01, div_ne01); \ \ float * dst_vtcm = (float *) dma_queue_pop(dma_queue).src; \ float * src0_vtcm = (float *) dma_queue_pop(dma_queue).dst; \ @@ -562,12 +597,12 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat \ const uint32_t next_ir = ir + block_size; \ if (next_ir < src0_end_row) { \ - const uint32_t next_block_size = unary_block_size(next_ir, src0_end_row, BLOCK, src0_contig, dst_contig,\ - ne01, div_ne01); \ + const uint32_t next_block_size = unary_block_size(next_ir, src0_end_row, BLOCK, block_src0_contig, \ + block_dst_contig, ne01, div_ne01); \ const uint32_t pref_ir = next_ir + next_block_size; \ if (pref_ir < src0_end_row) { \ - const uint32_t pref_block_size = unary_block_size(pref_ir, src0_end_row, BLOCK, src0_contig, \ - dst_contig, ne01, div_ne01); \ + const uint32_t pref_block_size = unary_block_size(pref_ir, src0_end_row, BLOCK, block_src0_contig, \ + block_dst_contig, ne01, div_ne01); \ const size_t src0_pref_off = src0_contig ? (pref_ir * nb01) : \ unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03); \ dma_queue_push(dma_queue, \ @@ -576,7 +611,8 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat \ if ((IS_RMS_NORM_MUL) && !uctx->broadcast_weight) { \ const size_t src1_pref_off = src1_contig ? (pref_ir * nb11) : \ - unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11, nb12, nb13); \ + unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11_bc, nb12_bc, \ + nb13_bc); \ dma_queue_push(dma_queue, \ dma_make_ptr(src1_vtcm, data_src1 + src1_pref_off), \ uctx->src1_row_size_aligned, nb11, uctx->src1_data_row_size, pref_block_size); \ @@ -603,6 +639,8 @@ DEFINE_UNARY_TASK(unary_silu, false, false, silu_f32(src0_vtcm, dst_vtcm, bl DEFINE_UNARY_TASK(unary_gelu, false, false, gelu_f32(src0_vtcm, dst_vtcm, block_size, uctx)) DEFINE_UNARY_TASK(unary_softplus, false, false, softplus_f32(src0_vtcm, dst_vtcm, block_size, uctx)) DEFINE_UNARY_TASK(unary_tanh, false, false, tanh_f32(src0_vtcm, dst_vtcm, block_size, uctx)) +DEFINE_UNARY_TASK(unary_abs, false, false, abs_f32(src0_vtcm, dst_vtcm, block_size, uctx)) +DEFINE_UNARY_TASK(unary_log, false, false, log_f32(src0_vtcm, dst_vtcm, block_size, uctx)) DEFINE_UNARY_TASK(l2_norm, false, false, l2_norm_f32(src0_vtcm, dst_vtcm, block_size, uctx)) DEFINE_UNARY_TASK(tri, false, true, tri_f32(src0_vtcm, dst_vtcm, block_size, ir, uctx)) @@ -850,6 +888,8 @@ DEFINE_UNARY_TILED_TASK(unary_silu, false, tile_silu_f32(dst_vtcm, src_vtcm, DEFINE_UNARY_TILED_TASK(unary_gelu, false, tile_gelu_f32(dst_vtcm, src_vtcm, tw)) DEFINE_UNARY_TILED_TASK(unary_softplus, false, tile_unary_softplus_f32(dst_vtcm, src_vtcm, tw)) DEFINE_UNARY_TILED_TASK(unary_tanh, false, hvx_tanh_f32_aa(dst_vtcm, src_vtcm, tw)) +DEFINE_UNARY_TILED_TASK(unary_abs, false, hvx_abs_f32_aa(dst_vtcm, src_vtcm, tw)) +DEFINE_UNARY_TILED_TASK(unary_log, false, hvx_log_f32_aa(dst_vtcm, src_vtcm, tw)) DEFINE_UNARY_TILED_TASK(tri, true, tri_apply_tile_f32(src_vtcm, dst_vtcm, tw, col, i01, ne0, tri_ttype)) static int execute_op_unary_f32(struct htp_ops_context * octx) { @@ -875,6 +915,8 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) { case HTP_OP_UNARY_GELU: op_type = "gelu-f32"; break; case HTP_OP_UNARY_SOFTPLUS: op_type = "softplus-f32"; break; case HTP_OP_UNARY_TANH: op_type = "tanh-f32"; break; + case HTP_OP_UNARY_ABS: op_type = "abs-f32"; break; + case HTP_OP_UNARY_LOG: op_type = "log-f32"; break; case HTP_OP_L2_NORM: op_type = "l2norm-f32"; break; case HTP_OP_TRI: op_type = "tri-f32"; break; @@ -973,6 +1015,8 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) { case HTP_OP_UNARY_GELU: task_func = unary_task_f32_tiled_unary_gelu; break; case HTP_OP_UNARY_SOFTPLUS: task_func = unary_task_f32_tiled_unary_softplus; break; case HTP_OP_UNARY_TANH: task_func = unary_task_f32_tiled_unary_tanh; break; + case HTP_OP_UNARY_ABS: task_func = unary_task_f32_tiled_unary_abs; break; + case HTP_OP_UNARY_LOG: task_func = unary_task_f32_tiled_unary_log; break; case HTP_OP_TRI: task_func = unary_task_f32_tiled_tri; break; default: break; } @@ -992,6 +1036,8 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) { case HTP_OP_UNARY_GELU: task_func = unary_task_f32_unary_gelu; break; case HTP_OP_UNARY_SOFTPLUS: task_func = unary_task_f32_unary_softplus; break; case HTP_OP_UNARY_TANH: task_func = unary_task_f32_unary_tanh; break; + case HTP_OP_UNARY_ABS: task_func = unary_task_f32_unary_abs; break; + case HTP_OP_UNARY_LOG: task_func = unary_task_f32_unary_log; break; case HTP_OP_L2_NORM: task_func = unary_task_f32_l2_norm; break; case HTP_OP_TRI: task_func = unary_task_f32_tri; break; default: break; diff --git a/ggml/src/ggml-hexagon/htp/unary-ops.h b/ggml/src/ggml-hexagon/htp/unary-ops.h index 1f4c3a5c4d9..458218ff443 100644 --- a/ggml/src/ggml-hexagon/htp/unary-ops.h +++ b/ggml/src/ggml-hexagon/htp/unary-ops.h @@ -55,6 +55,8 @@ static inline bool htp_op_is_unary(uint32_t opcode) { case HTP_OP_UNARY_GELU: case HTP_OP_UNARY_SOFTPLUS: case HTP_OP_UNARY_TANH: + case HTP_OP_UNARY_ABS: + case HTP_OP_UNARY_LOG: case HTP_OP_L2_NORM: case HTP_OP_TRI: return true; diff --git a/ggml/src/ggml-metal/ggml-metal-context.m b/ggml/src/ggml-metal/ggml-metal-context.m index 32d97cd5d0a..1227ed39a09 100644 --- a/ggml/src/ggml-metal/ggml-metal-context.m +++ b/ggml/src/ggml-metal/ggml-metal-context.m @@ -84,106 +84,108 @@ ggml_metal_t ggml_metal_init(ggml_metal_device_t dev) { GGML_LOG_INFO("%s: allocating\n", __func__); + @autoreleasepool { #if TARGET_OS_OSX && !GGML_METAL_NDEBUG - // Show all the Metal device instances in the system - NSArray * devices = MTLCopyAllDevices(); - for (id device in devices) { - GGML_LOG_INFO("%s: found device: %s\n", __func__, [[device name] UTF8String]); - } - [devices release]; // since it was created by a *Copy* C method + // Show all the Metal device instances in the system + NSArray * devices = MTLCopyAllDevices(); + for (id device in devices) { + GGML_LOG_INFO("%s: found device: %s\n", __func__, [[device name] UTF8String]); + } + [devices release]; // since it was created by a *Copy* C method #endif - // init context - ggml_metal_t res = calloc(1, sizeof(struct ggml_metal)); - - id device = ggml_metal_device_get_obj(dev); + // init context + ggml_metal_t res = calloc(1, sizeof(struct ggml_metal)); - GGML_LOG_INFO("%s: picking default device: %s\n", __func__, [[device name] UTF8String]); + id device = ggml_metal_device_get_obj(dev); - // TODO: would it be better to have one queue for the backend and one queue for the device? - // the graph encoders and async ops would use the backend queue while the sync ops would use the device queue? - //res->queue = [device newCommandQueue]; [TAG_QUEUE_PER_BACKEND] - id queue = ggml_metal_device_get_queue(dev); - if (queue == nil) { - GGML_LOG_ERROR("%s: error: failed to create command queue\n", __func__); - return NULL; - } + GGML_LOG_INFO("%s: picking default device: %s\n", __func__, [[device name] UTF8String]); - res->dev = dev; - res->lib = ggml_metal_device_get_library(dev); - if (res->lib == NULL) { - GGML_LOG_WARN("%s: the device does not have a precompiled Metal library - this is unexpected\n", __func__); - GGML_LOG_WARN("%s: will try to compile it on the fly\n", __func__); + // TODO: would it be better to have one queue for the backend and one queue for the device? + // the graph encoders and async ops would use the backend queue while the sync ops would use the device queue? + //res->queue = [device newCommandQueue]; [TAG_QUEUE_PER_BACKEND] + id queue = ggml_metal_device_get_queue(dev); + if (queue == nil) { + GGML_LOG_ERROR("%s: error: failed to create command queue\n", __func__); + return NULL; + } - res->lib = ggml_metal_library_init(dev); + res->dev = dev; + res->lib = ggml_metal_device_get_library(dev); if (res->lib == NULL) { - GGML_LOG_ERROR("%s: error: failed to initialize the Metal library\n", __func__); + GGML_LOG_WARN("%s: the device does not have a precompiled Metal library - this is unexpected\n", __func__); + GGML_LOG_WARN("%s: will try to compile it on the fly\n", __func__); - free(res); + res->lib = ggml_metal_library_init(dev); + if (res->lib == NULL) { + GGML_LOG_ERROR("%s: error: failed to initialize the Metal library\n", __func__); - return NULL; + free(res); + + return NULL; + } } - } - res->ev_cpy = ggml_metal_device_event_init(dev); + res->ev_cpy = ggml_metal_device_event_init(dev); - const struct ggml_metal_device_props * props_dev = ggml_metal_device_get_props(dev); + const struct ggml_metal_device_props * props_dev = ggml_metal_device_get_props(dev); - snprintf(res->name, sizeof(res->name), "%s", props_dev->name); + snprintf(res->name, sizeof(res->name), "%s", props_dev->name); - res->d_queue = dispatch_queue_create("ggml-metal", DISPATCH_QUEUE_CONCURRENT); + res->d_queue = dispatch_queue_create("ggml-metal", DISPATCH_QUEUE_CONCURRENT); - res->use_fusion = getenv("GGML_METAL_FUSION_DISABLE") == nil; - res->use_concurrency = getenv("GGML_METAL_CONCURRENCY_DISABLE") == nil; + res->use_fusion = getenv("GGML_METAL_FUSION_DISABLE") == nil; + res->use_concurrency = getenv("GGML_METAL_CONCURRENCY_DISABLE") == nil; - { - const char * val = getenv("GGML_METAL_GRAPH_DEBUG"); - res->debug_graph = val ? atoi(val) : 0; - } + { + const char * val = getenv("GGML_METAL_GRAPH_DEBUG"); + res->debug_graph = val ? atoi(val) : 0; + } - { - const char * val = getenv("GGML_METAL_FUSION_DEBUG"); - res->debug_fusion = val ? atoi(val) : 0; - } + { + const char * val = getenv("GGML_METAL_FUSION_DEBUG"); + res->debug_fusion = val ? atoi(val) : 0; + } - res->use_graph_optimize = true; + res->use_graph_optimize = true; - if (getenv("GGML_METAL_GRAPH_OPTIMIZE_DISABLE") != NULL) { - res->use_graph_optimize = false; - } + if (getenv("GGML_METAL_GRAPH_OPTIMIZE_DISABLE") != NULL) { + res->use_graph_optimize = false; + } - memset(res->fuse_cnt, 0, sizeof(res->fuse_cnt)); + memset(res->fuse_cnt, 0, sizeof(res->fuse_cnt)); - GGML_LOG_INFO("%s: use fusion = %s\n", __func__, res->use_fusion ? "true" : "false"); - GGML_LOG_INFO("%s: use concurrency = %s\n", __func__, res->use_concurrency ? "true" : "false"); - GGML_LOG_INFO("%s: use graph optimize = %s\n", __func__, res->use_graph_optimize ? "true" : "false"); + GGML_LOG_INFO("%s: use fusion = %s\n", __func__, res->use_fusion ? "true" : "false"); + GGML_LOG_INFO("%s: use concurrency = %s\n", __func__, res->use_concurrency ? "true" : "false"); + GGML_LOG_INFO("%s: use graph optimize = %s\n", __func__, res->use_graph_optimize ? "true" : "false"); - res->capture_compute = 0; - res->capture_started = false; - res->capture_scope = nil; + res->capture_compute = 0; + res->capture_started = false; + res->capture_scope = nil; - { - const char * val = getenv("GGML_METAL_CAPTURE_COMPUTE"); - if (val) { - res->capture_compute = atoi(val); + { + const char * val = getenv("GGML_METAL_CAPTURE_COMPUTE"); + if (val) { + res->capture_compute = atoi(val); + } } - } - res->has_error = false; + res->has_error = false; - res->gf = nil; - res->encode_async = nil; - for (int i = 0; i < GGML_METAL_MAX_COMMAND_BUFFERS; ++i) { - res->cmd_bufs[i].obj = nil; - } + res->gf = nil; + res->encode_async = nil; + for (int i = 0; i < GGML_METAL_MAX_COMMAND_BUFFERS; ++i) { + res->cmd_bufs[i].obj = nil; + } - res->cmd_bufs_ext = [[NSMutableArray alloc] init]; + res->cmd_bufs_ext = [[NSMutableArray alloc] init]; - res->cmd_buf_last = nil; + res->cmd_buf_last = nil; - res->pipelines_ext = ggml_metal_pipelines_init(); + res->pipelines_ext = ggml_metal_pipelines_init(); - return res; + return res; + } } void ggml_metal_free(ggml_metal_t ctx) { diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index a82caa5e430..4e855be4467 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -593,7 +593,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_me // - sgptg floats for shared_x_dt (nsg) // - sgptg floats for shared_dA (nsg) // Total: nsg * (32 + 2) floats - res.smem = (32 + 2)*sizeof(float)*nsg; + res.smem = GGML_PAD((32 + 2)*sizeof(float)*nsg, 16); return res; } @@ -1029,6 +1029,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0(g } res.smem = (size_t) ne02*ne20*sizeof(uint16_t); + res.smem = GGML_PAD(res.smem, 16); return res; } diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 19c57820e85..85c0f576b9f 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -778,7 +778,9 @@ void ggml_metal_encoder_free(ggml_metal_encoder_t encoder) { } void ggml_metal_encoder_debug_group_push(ggml_metal_encoder_t encoder, const char * name) { - [encoder->obj pushDebugGroup:[NSString stringWithCString:name encoding:NSUTF8StringEncoding]]; + @autoreleasepool { + [encoder->obj pushDebugGroup:[NSString stringWithCString:name encoding:NSUTF8StringEncoding]]; + } } void ggml_metal_encoder_debug_group_pop (ggml_metal_encoder_t encoder) { @@ -798,6 +800,9 @@ void ggml_metal_encoder_set_buffer(ggml_metal_encoder_t encoder, struct ggml_met } void ggml_metal_encoder_set_threadgroup_memory_size(ggml_metal_encoder_t encoder, size_t size, int idx) { + // ref: https://developer.apple.com/documentation/metal/mtlcomputecommandencoder/setthreadgroupmemorylength(_:index:) + GGML_ASSERT(size % 16 == 0); + [encoder->obj setThreadgroupMemoryLength:size atIndex:idx]; } @@ -1023,249 +1028,251 @@ ggml_metal_device_t ggml_metal_device_init(int device, int n_devices) { assert(dev != NULL); - if (dev->mtl_device == nil) { - dev->mtl_device = MTLCreateSystemDefaultDevice(); - - if (dev->mtl_device) { - dev->mtl_queue = [dev->mtl_device newCommandQueue]; - if (dev->mtl_queue == nil) { - GGML_LOG_ERROR("%s: error: failed to create command queue\n", __func__); - } + @autoreleasepool { + if (dev->mtl_device == nil) { + dev->mtl_device = MTLCreateSystemDefaultDevice(); - dev->addr_virt = 0x000000400ULL; + if (dev->mtl_device) { + dev->mtl_queue = [dev->mtl_device newCommandQueue]; + if (dev->mtl_queue == nil) { + GGML_LOG_ERROR("%s: error: failed to create command queue\n", __func__); + } - dev->props.device = device; + dev->addr_virt = 0x000000400ULL; - // the Metal backend uses the system default device as the single physical device; - // additional (virtual) devices are emulated on top of it via GGML_METAL_DEVICES - dev->props.device_phys = 0; - dev->props.device_virt = device; + dev->props.device = device; - dev->props.has_simdgroup_reduction = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7]; - dev->props.has_simdgroup_reduction |= [dev->mtl_device supportsFamily:MTLGPUFamilyMetal3_GGML]; + // the Metal backend uses the system default device as the single physical device; + // additional (virtual) devices are emulated on top of it via GGML_METAL_DEVICES + dev->props.device_phys = 0; + dev->props.device_virt = device; - dev->props.has_simdgroup_mm = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7]; - dev->props.has_unified_memory = dev->mtl_device.hasUnifiedMemory; + dev->props.has_simdgroup_reduction = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7]; + dev->props.has_simdgroup_reduction |= [dev->mtl_device supportsFamily:MTLGPUFamilyMetal3_GGML]; - dev->props.has_bfloat = [dev->mtl_device supportsFamily:MTLGPUFamilyMetal3_GGML]; - dev->props.has_bfloat |= [dev->mtl_device supportsFamily:MTLGPUFamilyApple6]; - if (getenv("GGML_METAL_BF16_DISABLE") != NULL) { - dev->props.has_bfloat = false; - } + dev->props.has_simdgroup_mm = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7]; + dev->props.has_unified_memory = dev->mtl_device.hasUnifiedMemory; - dev->props.has_tensor = [dev->mtl_device supportsFamily:MTLGPUFamilyMetal4_GGML]; - if (getenv("GGML_METAL_TENSOR_DISABLE") != NULL) { - dev->props.has_tensor = false; - } + dev->props.has_bfloat = [dev->mtl_device supportsFamily:MTLGPUFamilyMetal3_GGML]; + dev->props.has_bfloat |= [dev->mtl_device supportsFamily:MTLGPUFamilyApple6]; + if (getenv("GGML_METAL_BF16_DISABLE") != NULL) { + dev->props.has_bfloat = false; + } - // note: disable the tensor API by default for old chips because with the current implementation it is not useful - // - M2 Ultra: ~5% slower - // - M4, M4 Max: no significant difference - // - // TODO: try to update the tensor API kernels to at least match the simdgroup performance - if (getenv("GGML_METAL_TENSOR_ENABLE") == NULL && - ![[dev->mtl_device name] containsString:@"M5"] && - ![[dev->mtl_device name] containsString:@"M6"] && - ![[dev->mtl_device name] containsString:@"A19"] && - ![[dev->mtl_device name] containsString:@"A20"]) { - GGML_LOG_INFO("%s: tensor API disabled for pre-M5 and pre-A19 devices\n", __func__); - dev->props.has_tensor = false; - } + dev->props.has_tensor = [dev->mtl_device supportsFamily:MTLGPUFamilyMetal4_GGML]; + if (getenv("GGML_METAL_TENSOR_DISABLE") != NULL) { + dev->props.has_tensor = false; + } - // double-check that the tensor API compiles - if (dev->props.has_tensor) { - const char * src_tensor_f16 = "\n" - "#include \n" - "#include \n" - "#include \n" - " \n" - "using namespace metal; \n" - "using namespace mpp::tensor_ops; \n" - " \n" - "kernel void dummy_kernel( \n" - " tensor> A [[buffer(0)]], \n" - " tensor> B [[buffer(1)]], \n" - " device float * C [[buffer(2)]], \n" - " uint2 tgid [[threadgroup_position_in_grid]]) \n" - "{ \n" - " auto tA = A.slice(0, (int)tgid.y); \n" - " auto tB = B.slice((int)tgid.x, 0); \n" - " \n" - " matmul2d< \n" - " matmul2d_descriptor(16, 16, dynamic_extent), \n" - " execution_simdgroups<4>> mm; \n" - " \n" - " auto cT = mm.get_destination_cooperative_tensor(); \n" - " \n" - " auto sA = tA.slice(0, 0); \n" - " auto sB = tB.slice(0, 0); \n" - " mm.run(sB, sA, cT); \n" - " \n" - " auto tC = tensor, tensor_inline>(C, dextents(16, 16)); \n" - " \n" - " cT.store(tC); \n" - "}"; - - GGML_LOG_INFO("%s: testing tensor API for f16 support\n", __func__); - ggml_metal_library_t lib = ggml_metal_library_init_from_source(dev, src_tensor_f16, false); - if (lib == NULL) { - GGML_LOG_WARN("%s: - the tensor API is not supported in this environment - disabling\n", __func__); + // note: disable the tensor API by default for old chips because with the current implementation it is not useful + // - M2 Ultra: ~5% slower + // - M4, M4 Max: no significant difference + // + // TODO: try to update the tensor API kernels to at least match the simdgroup performance + if (getenv("GGML_METAL_TENSOR_ENABLE") == NULL && + ![[dev->mtl_device name] containsString:@"M5"] && + ![[dev->mtl_device name] containsString:@"M6"] && + ![[dev->mtl_device name] containsString:@"A19"] && + ![[dev->mtl_device name] containsString:@"A20"]) { + GGML_LOG_INFO("%s: tensor API disabled for pre-M5 and pre-A19 devices\n", __func__); dev->props.has_tensor = false; - } else { - struct ggml_metal_pipeline_with_params ppl = ggml_metal_library_compile_pipeline(lib, "dummy_kernel", "dummy_kernel", nil); - if (!ppl.pipeline) { + } + + // double-check that the tensor API compiles + if (dev->props.has_tensor) { + const char * src_tensor_f16 = "\n" + "#include \n" + "#include \n" + "#include \n" + " \n" + "using namespace metal; \n" + "using namespace mpp::tensor_ops; \n" + " \n" + "kernel void dummy_kernel( \n" + " tensor> A [[buffer(0)]], \n" + " tensor> B [[buffer(1)]], \n" + " device float * C [[buffer(2)]], \n" + " uint2 tgid [[threadgroup_position_in_grid]]) \n" + "{ \n" + " auto tA = A.slice(0, (int)tgid.y); \n" + " auto tB = B.slice((int)tgid.x, 0); \n" + " \n" + " matmul2d< \n" + " matmul2d_descriptor(16, 16, dynamic_extent), \n" + " execution_simdgroups<4>> mm; \n" + " \n" + " auto cT = mm.get_destination_cooperative_tensor(); \n" + " \n" + " auto sA = tA.slice(0, 0); \n" + " auto sB = tB.slice(0, 0); \n" + " mm.run(sB, sA, cT); \n" + " \n" + " auto tC = tensor, tensor_inline>(C, dextents(16, 16)); \n" + " \n" + " cT.store(tC); \n" + "}"; + + GGML_LOG_INFO("%s: testing tensor API for f16 support\n", __func__); + ggml_metal_library_t lib = ggml_metal_library_init_from_source(dev, src_tensor_f16, false); + if (lib == NULL) { GGML_LOG_WARN("%s: - the tensor API is not supported in this environment - disabling\n", __func__); dev->props.has_tensor = false; - } + } else { + struct ggml_metal_pipeline_with_params ppl = ggml_metal_library_compile_pipeline(lib, "dummy_kernel", "dummy_kernel", nil); + if (!ppl.pipeline) { + GGML_LOG_WARN("%s: - the tensor API is not supported in this environment - disabling\n", __func__); + dev->props.has_tensor = false; + } - ggml_metal_library_free(lib); + ggml_metal_library_free(lib); + } } - } - // try to compile a dummy kernel to determine if the tensor API is supported for bfloat - if (dev->props.has_tensor && dev->props.has_bfloat) { - const char * src_tensor_bf16 = "\n" - "#include \n" - "#include \n" - "#include \n" - " \n" - "using namespace metal; \n" - "using namespace mpp::tensor_ops; \n" - " \n" - "kernel void dummy_kernel( \n" - " tensor> A [[buffer(0)]], \n" - " tensor> B [[buffer(1)]], \n" - " device float * C [[buffer(2)]], \n" - " uint2 tgid [[threadgroup_position_in_grid]]) \n" - "{ \n" - " auto tA = A.slice(0, (int)tgid.y); \n" - " auto tB = B.slice((int)tgid.x, 0); \n" - " \n" - " matmul2d< \n" - " matmul2d_descriptor(16, 16, dynamic_extent), \n" - " execution_simdgroups<4>> mm; \n" - " \n" - " auto cT = mm.get_destination_cooperative_tensor(); \n" - " \n" - " auto sA = tA.slice(0, 0); \n" - " auto sB = tB.slice(0, 0); \n" - " mm.run(sB, sA, cT); \n" - " \n" - " auto tC = tensor, tensor_inline>(C, dextents(16, 16)); \n" - " \n" - " cT.store(tC); \n" - "}"; - - GGML_LOG_INFO("%s: testing tensor API for bfloat support\n", __func__); - ggml_metal_library_t lib = ggml_metal_library_init_from_source(dev, src_tensor_bf16, false); - if (lib == NULL) { - GGML_LOG_WARN("%s: - the tensor API does not support bfloat - disabling bfloat support\n", __func__); - dev->props.has_bfloat = false; - } else { - struct ggml_metal_pipeline_with_params ppl = ggml_metal_library_compile_pipeline(lib, "dummy_kernel", "dummy_kernel", nil); - if (!ppl.pipeline) { + // try to compile a dummy kernel to determine if the tensor API is supported for bfloat + if (dev->props.has_tensor && dev->props.has_bfloat) { + const char * src_tensor_bf16 = "\n" + "#include \n" + "#include \n" + "#include \n" + " \n" + "using namespace metal; \n" + "using namespace mpp::tensor_ops; \n" + " \n" + "kernel void dummy_kernel( \n" + " tensor> A [[buffer(0)]], \n" + " tensor> B [[buffer(1)]], \n" + " device float * C [[buffer(2)]], \n" + " uint2 tgid [[threadgroup_position_in_grid]]) \n" + "{ \n" + " auto tA = A.slice(0, (int)tgid.y); \n" + " auto tB = B.slice((int)tgid.x, 0); \n" + " \n" + " matmul2d< \n" + " matmul2d_descriptor(16, 16, dynamic_extent), \n" + " execution_simdgroups<4>> mm; \n" + " \n" + " auto cT = mm.get_destination_cooperative_tensor(); \n" + " \n" + " auto sA = tA.slice(0, 0); \n" + " auto sB = tB.slice(0, 0); \n" + " mm.run(sB, sA, cT); \n" + " \n" + " auto tC = tensor, tensor_inline>(C, dextents(16, 16)); \n" + " \n" + " cT.store(tC); \n" + "}"; + + GGML_LOG_INFO("%s: testing tensor API for bfloat support\n", __func__); + ggml_metal_library_t lib = ggml_metal_library_init_from_source(dev, src_tensor_bf16, false); + if (lib == NULL) { GGML_LOG_WARN("%s: - the tensor API does not support bfloat - disabling bfloat support\n", __func__); dev->props.has_bfloat = false; - } + } else { + struct ggml_metal_pipeline_with_params ppl = ggml_metal_library_compile_pipeline(lib, "dummy_kernel", "dummy_kernel", nil); + if (!ppl.pipeline) { + GGML_LOG_WARN("%s: - the tensor API does not support bfloat - disabling bfloat support\n", __func__); + dev->props.has_bfloat = false; + } - ggml_metal_library_free(lib); + ggml_metal_library_free(lib); + } } - } - dev->props.use_residency_sets = true; + dev->props.use_residency_sets = true; #if defined(GGML_METAL_HAS_RESIDENCY_SETS) - dev->props.use_residency_sets = getenv("GGML_METAL_NO_RESIDENCY") == nil; + dev->props.use_residency_sets = getenv("GGML_METAL_NO_RESIDENCY") == nil; #endif - dev->props.use_shared_buffers = dev->props.has_unified_memory; + dev->props.use_shared_buffers = dev->props.has_unified_memory; #if TARGET_OS_OSX - // In case of eGPU, shared memory may be preferable. - dev->props.use_shared_buffers |= [dev->mtl_device location] == MTLDeviceLocationExternal; + // In case of eGPU, shared memory may be preferable. + dev->props.use_shared_buffers |= [dev->mtl_device location] == MTLDeviceLocationExternal; #endif - if (getenv("GGML_METAL_SHARED_BUFFERS_DISABLE") != NULL) { - dev->props.use_shared_buffers = false; - } - if (getenv("GGML_METAL_SHARED_BUFFERS_ENABLE") != NULL) { - dev->props.use_shared_buffers = true; - } + if (getenv("GGML_METAL_SHARED_BUFFERS_DISABLE") != NULL) { + dev->props.use_shared_buffers = false; + } + if (getenv("GGML_METAL_SHARED_BUFFERS_ENABLE") != NULL) { + dev->props.use_shared_buffers = true; + } - dev->props.supports_gpu_family_apple7 = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7]; + dev->props.supports_gpu_family_apple7 = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7]; - dev->props.device_id = ggml_metal_device_id_parse([[dev->mtl_device name] UTF8String]); + dev->props.device_id = ggml_metal_device_id_parse([[dev->mtl_device name] UTF8String]); - dev->props.op_offload_min_batch_size = getenv("GGML_OP_OFFLOAD_MIN_BATCH") ? atoi(getenv("GGML_OP_OFFLOAD_MIN_BATCH")) : 32; + dev->props.op_offload_min_batch_size = getenv("GGML_OP_OFFLOAD_MIN_BATCH") ? atoi(getenv("GGML_OP_OFFLOAD_MIN_BATCH")) : 32; - dev->props.max_buffer_size = dev->mtl_device.maxBufferLength; - dev->props.max_theadgroup_memory_size = dev->mtl_device.maxThreadgroupMemoryLength; - if (@available(macOS 10.12, iOS 16.0, *)) { - dev->props.max_working_set_size = dev->mtl_device.recommendedMaxWorkingSetSize; - } else { - dev->props.max_working_set_size = dev->mtl_device.maxBufferLength; - } - - snprintf(dev->props.name, sizeof(dev->props.name), "%s%d", "MTL", device); - const char * gpu_name = [[dev->mtl_device name] UTF8String]; - if (n_devices > 1) { - snprintf(dev->props.desc, sizeof(dev->props.desc), "%s (dev p%d/v%d)", - gpu_name, dev->props.device_phys, dev->props.device_virt); - } else { - snprintf(dev->props.desc, sizeof(dev->props.desc), "%s", gpu_name); - } + dev->props.max_buffer_size = dev->mtl_device.maxBufferLength; + dev->props.max_theadgroup_memory_size = dev->mtl_device.maxThreadgroupMemoryLength; + if (@available(macOS 10.12, iOS 16.0, *)) { + dev->props.max_working_set_size = dev->mtl_device.recommendedMaxWorkingSetSize; + } else { + dev->props.max_working_set_size = dev->mtl_device.maxBufferLength; + } - dev->library = ggml_metal_library_init(dev); - if (!dev->library) { - GGML_LOG_ERROR("%s: error: failed to create library\n", __func__); - } + snprintf(dev->props.name, sizeof(dev->props.name), "%s%d", "MTL", device); + const char * gpu_name = [[dev->mtl_device name] UTF8String]; + if (n_devices > 1) { + snprintf(dev->props.desc, sizeof(dev->props.desc), "%s (dev p%d/v%d)", + gpu_name, dev->props.device_phys, dev->props.device_virt); + } else { + snprintf(dev->props.desc, sizeof(dev->props.desc), "%s", gpu_name); + } - if (dev->props.use_residency_sets) { - dev->rsets = ggml_metal_rsets_init(dev); - } else { - dev->rsets = nil; - } + dev->library = ggml_metal_library_init(dev); + if (!dev->library) { + GGML_LOG_ERROR("%s: error: failed to create library\n", __func__); + } - // print MTL GPU family: - GGML_LOG_INFO("%s: GPU name: %s (%s)\n", __func__, dev->props.name, dev->props.desc); + if (dev->props.use_residency_sets) { + dev->rsets = ggml_metal_rsets_init(dev); + } else { + dev->rsets = nil; + } - // determine max supported GPU family - // https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf - // https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf - { - for (int i = MTLGPUFamilyApple1 + 20; i >= MTLGPUFamilyApple1; --i) { - if ([dev->mtl_device supportsFamily:i]) { - dev->props.gpu_family = i - (int) MTLGPUFamilyApple1 + 1; - GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyApple%d (%d)\n", __func__, dev->props.gpu_family, i); - break; + // print MTL GPU family: + GGML_LOG_INFO("%s: GPU name: %s (%s)\n", __func__, dev->props.name, dev->props.desc); + + // determine max supported GPU family + // https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf + // https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf + { + for (int i = MTLGPUFamilyApple1 + 20; i >= MTLGPUFamilyApple1; --i) { + if ([dev->mtl_device supportsFamily:i]) { + dev->props.gpu_family = i - (int) MTLGPUFamilyApple1 + 1; + GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyApple%d (%d)\n", __func__, dev->props.gpu_family, i); + break; + } } - } - for (int i = MTLGPUFamilyCommon1 + 5; i >= MTLGPUFamilyCommon1; --i) { - if ([dev->mtl_device supportsFamily:i]) { - GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyCommon%d (%d)\n", __func__, i - (int) MTLGPUFamilyCommon1 + 1, i); - break; + for (int i = MTLGPUFamilyCommon1 + 5; i >= MTLGPUFamilyCommon1; --i) { + if ([dev->mtl_device supportsFamily:i]) { + GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyCommon%d (%d)\n", __func__, i - (int) MTLGPUFamilyCommon1 + 1, i); + break; + } } - } - for (int i = MTLGPUFamilyMetal3_GGML + 5; i >= MTLGPUFamilyMetal3_GGML; --i) { - if ([dev->mtl_device supportsFamily:i]) { - GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyMetal%d (%d)\n", __func__, i - (int) MTLGPUFamilyMetal3_GGML + 3, i); - break; + for (int i = MTLGPUFamilyMetal3_GGML + 5; i >= MTLGPUFamilyMetal3_GGML; --i) { + if ([dev->mtl_device supportsFamily:i]) { + GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyMetal%d (%d)\n", __func__, i - (int) MTLGPUFamilyMetal3_GGML + 3, i); + break; + } } } - } - GGML_LOG_INFO("%s: simdgroup reduction = %s\n", __func__, dev->props.has_simdgroup_reduction ? "true" : "false"); - GGML_LOG_INFO("%s: simdgroup matrix mul. = %s\n", __func__, dev->props.has_simdgroup_mm ? "true" : "false"); - GGML_LOG_INFO("%s: has unified memory = %s\n", __func__, dev->props.has_unified_memory ? "true" : "false"); - GGML_LOG_INFO("%s: has bfloat = %s\n", __func__, dev->props.has_bfloat ? "true" : "false"); - GGML_LOG_INFO("%s: has tensor = %s\n", __func__, dev->props.has_tensor ? "true" : "false"); - GGML_LOG_INFO("%s: use residency sets = %s\n", __func__, dev->props.use_residency_sets ? "true" : "false"); - GGML_LOG_INFO("%s: use shared buffers = %s\n", __func__, dev->props.use_shared_buffers ? "true" : "false"); + GGML_LOG_INFO("%s: simdgroup reduction = %s\n", __func__, dev->props.has_simdgroup_reduction ? "true" : "false"); + GGML_LOG_INFO("%s: simdgroup matrix mul. = %s\n", __func__, dev->props.has_simdgroup_mm ? "true" : "false"); + GGML_LOG_INFO("%s: has unified memory = %s\n", __func__, dev->props.has_unified_memory ? "true" : "false"); + GGML_LOG_INFO("%s: has bfloat = %s\n", __func__, dev->props.has_bfloat ? "true" : "false"); + GGML_LOG_INFO("%s: has tensor = %s\n", __func__, dev->props.has_tensor ? "true" : "false"); + GGML_LOG_INFO("%s: use residency sets = %s\n", __func__, dev->props.use_residency_sets ? "true" : "false"); + GGML_LOG_INFO("%s: use shared buffers = %s\n", __func__, dev->props.use_shared_buffers ? "true" : "false"); #if TARGET_OS_OSX || (TARGET_OS_IOS && __clang_major__ >= 15) - if (@available(macOS 10.12, iOS 16.0, *)) { - GGML_LOG_INFO("%s: recommendedMaxWorkingSetSize = %8.2f MB\n", __func__, dev->props.max_working_set_size / 1e6); - } + if (@available(macOS 10.12, iOS 16.0, *)) { + GGML_LOG_INFO("%s: recommendedMaxWorkingSetSize = %8.2f MB\n", __func__, dev->props.max_working_set_size / 1e6); + } #endif + } } } diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 9becf04797b..49102afe9c0 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -660,6 +660,7 @@ typedef struct { uint64_t nb0; uint64_t nb1; uint64_t nb2; + uint64_t nb3; } ggml_metal_kargs_conv_transpose_2d; typedef struct { diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 75de0f6dd08..89c8483b371 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -948,7 +948,7 @@ int ggml_metal_op_sum(ggml_metal_op_t ctx, int idx) { ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); - ggml_metal_encoder_set_threadgroup_memory_size(enc, nsg * sizeof(float), 0); + ggml_metal_encoder_set_threadgroup_memory_size(enc, GGML_PAD(nsg * sizeof(float), 16), 0); ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, nth, 1, 1); @@ -4645,6 +4645,7 @@ int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) { const int32_t OW = op->ne[0]; const int32_t OH = op->ne[1]; const int32_t OC = op->ne[2]; + const int32_t N = op->src[1]->ne[3]; ggml_metal_kargs_conv_transpose_2d args = { /*.IC =*/ IC, @@ -4657,6 +4658,7 @@ int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) { /*.nb0 =*/ nb0, /*.nb1 =*/ nb1, /*.nb2 =*/ nb2, + /*.nb3 =*/ nb3, }; auto pipeline = ggml_metal_library_get_pipeline_conv_transpose_2d(lib, op); @@ -4671,7 +4673,7 @@ int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) { const size_t smem = GGML_PAD(KW * KH * sizeof(float), 16); ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); - ggml_metal_encoder_dispatch_threadgroups(enc, OW, OH, OC, KW, KH, 1); + ggml_metal_encoder_dispatch_threadgroups(enc, OW, OH, OC * N, KW, KH, 1); return 1; } diff --git a/ggml/src/ggml-metal/ggml-metal-tuning.cpp b/ggml/src/ggml-metal/ggml-metal-tuning.cpp index 6d8c18e6a6a..26b634d2143 100644 --- a/ggml/src/ggml-metal/ggml-metal-tuning.cpp +++ b/ggml/src/ggml-metal/ggml-metal-tuning.cpp @@ -66,6 +66,7 @@ fa_vec_cfg_t fa_vec_baseline_cfg(int dk, int dv) { // One row per kept bucket, plus per-(dtype,dk,dv) ne11-collapsed domain defaults // (ne11_b = FA_VEC_NE11_DEFAULT, ne01_b = domain). To retune or add a device, re-run the // sweep and paste its output. See ggml-metal-tuning.h for the row/lookup semantics. +// ref: https://github.com/ggml-org/llama.cpp/pull/27824 constexpr fa_vec_entry_t fa_vec_tuned_table[] = { { { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_F16, 32, 32, 3, 3 }, { 2, 4 } }, { { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } }, @@ -278,6 +279,194 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = { { { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M1_PRO, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 64, 64, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 128, 128, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 128, 128, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 128, 128, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 192, 128, 1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 192, 128, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 320, 256, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 320, 256, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_F16, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 32, 32, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 32, 32, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 32, 32, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 64, 64, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 64, 64, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 96, 96, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 96, 96, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 96, 96, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 96, 96, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 320, 256, 1, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 32, 32, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 32, 32, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 32, 32, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 32, 32, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 96, 96, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 96, 96, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 128, 128, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 128, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 320, 256, 1, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q4_1, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 128, 128, 1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 128, 128, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 128, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 192, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 192, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 192, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 128, 1, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 128, 2, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 192, 128, 3, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 256, 256, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 320, 256, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 128, 128, 2, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 128, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 192, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 128, 1, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 128, 2, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 192, 128, 3, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 256, 256, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 256, 256, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 320, 256, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q5_1, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 32, 32, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 32, 32, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 32, 32, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 64, 64, 1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 64, 64, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 96, 96, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 96, 96, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 96, 96, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 320, 256, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 320, 256, 1, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 320, 256, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 320, 256, 2, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 320, 256, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 320, 256, 3, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M1_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_F16, 64, 64, -1, 1 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_F16, 128, 128, 2, 1 }, { 1, 4 } }, @@ -449,6 +638,535 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = { { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M2_ULTRA, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 1, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, 3, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, 3, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, 3, 3 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 96, 96, 3, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 128, 128, 2, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 192, 1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 192, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 128, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 128, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 256, 256, 2, 3 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 320, 256, 1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 320, 256, 3, 0 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 2, 0 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 3, 0 }, { 4, 1 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 1, 3 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 2, 1 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 2, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 3, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 512, 512, 3, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 576, 512, 2, 0 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 576, 512, 2, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 576, 512, 2, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_F16, 576, 512, 3, 1 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 64, 64, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 96, 96, 3, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, 1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 128, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 192, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, 1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M3_MAX, GGML_TYPE_Q8_0, 576, 512, 1, 4 }, { 1, 2 } }, + + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 32, 32, 1, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 32, 32, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 64, 64, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 96, 96, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 96, 96, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 128, 128, 1, 2 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 128, 128, 2, 2 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 128, 128, 2, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 192, 192, 3, 0 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 192, 192, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 192, 128, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 192, 128, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 256, 256, 3, 0 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 320, 256, 3, 0 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 512, 512, 3, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 512, 512, 3, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 512, 512, 3, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_F16, 512, 512, 3, 3 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 64, 64, 1, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 64, 64, 2, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 64, 64, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 64, 64, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 128, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 192, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 576, 512, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 576, 512, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 576, 512, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 576, 512, 1, 3 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_0, 576, 512, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 64, 64, 1, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 64, 64, 2, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 64, 64, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 64, 64, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 128, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 192, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 576, 512, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 576, 512, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 576, 512, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 576, 512, 1, 3 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q4_1, 576, 512, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 96, 96, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 96, 96, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, 1, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, 1, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, 2, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 128, 128, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 256, 256, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 256, 256, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 320, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 320, 256, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 320, 256, 2, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 320, 256, 3, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 576, 512, 2, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 576, 512, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 576, 512, 2, 3 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_0, 576, 512, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 96, 96, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 96, 96, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 128, 128, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 128, 128, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 128, 128, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 128, 128, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 128, 128, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 128, 128, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 256, 256, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 256, 256, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 320, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 320, 256, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 320, 256, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 320, 256, 2, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 320, 256, 3, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 576, 512, 2, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 576, 512, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 576, 512, 2, 3 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q5_1, 576, 512, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 64, 64, 1, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 64, 64, 2, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 64, 64, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 192, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 192, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 192, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 192, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 576, 512, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 32, 32, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 64, 64, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 96, 96, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 96, 96, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 128, 128, 1, 2 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 192, 128, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 256, 256, 3, 0 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 320, 256, 3, 0 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_F16, 512, 512, 3, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 32, 32, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 64, 64, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 128, 128, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 128, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 576, 512, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 576, 512, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 576, 512, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 576, 512, 1, 3 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_0, 576, 512, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 32, 32, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 96, 96, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 96, 96, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 128, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 192, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 512, 512, 3, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q4_1, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 96, 96, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 96, 96, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 128, 128, 2, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 128, 128, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 256, 256, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 256, 256, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 256, 256, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 256, 256, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 320, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 320, 256, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 320, 256, 2, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 320, 256, 2, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 320, 256, 3, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 320, 256, 3, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 576, 512, 2, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 576, 512, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 576, 512, 2, 3 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_0, 576, 512, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 96, 96, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 96, 96, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 96, 96, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 96, 96, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 128, 128, 2, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 128, 128, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 128, 128, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 128, 128, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 192, 192, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 256, 256, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 256, 256, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 256, 256, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 320, 256, 1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 320, 256, 1, 3 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 320, 256, 1, 4 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 320, 256, 3, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 576, 512, 2, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 576, 512, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 576, 512, 2, 3 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q5_1, 576, 512, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 128, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 1, 3 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_PRO, GGML_TYPE_Q8_0, 576, 512, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } }, { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 32, 32, 2, 3 }, { 2, 4 } }, { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } }, @@ -640,7 +1358,215 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = { { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M4_MAX, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, - { { GGML_METAL_DEVICE_M5_MAX, GGML_TYPE_F16, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 1, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 32, 32, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 96, 96, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 96, 96, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 3, 0 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 1, 2 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 3, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 128, 128, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 192, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 192, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 3, 2 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 192, 128, 3, 3 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 256, 256, 1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 256, 256, 1, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 320, 256, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 320, 256, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 320, 256, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 320, 256, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 2, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 2, 3 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 3, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 3, 2 }, { 4, 1 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 3, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 512, 512, 3, 4 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 576, 512, 2, 0 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 576, 512, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 576, 512, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 576, 512, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_F16, 576, 512, 2, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 32, 32, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, 1, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, 1, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, 2, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 64, 64, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 96, 96, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 96, 96, 3, 4 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 192, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 192, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 192, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 128, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 320, 256, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 512, 512, 3, 0 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 512, 512, 2, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 512, 512, 3, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 64, 64, 1, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 64, 64, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 64, 64, 3, 4 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 128, 128, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 128, 128, 1, 4 }, { 1, 1 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 192, 192, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 192, 192, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 192, 128, 2, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 192, 128, 3, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 256, 256, -1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 256, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 256, 256, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 320, 256, 3, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 320, 256, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 320, 256, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 512, 512, 2, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 512, 512, 2, 3 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 512, 512, 3, 3 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, -1, 1 }, { 2, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, 1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, 1, 3 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_F16, 576, 512, 2, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 32, 32, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 64, 64, 1, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 96, 96, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 192, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 192, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 512, 512, 1, 0 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 512, 512, 2, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q4_0, 512, 512, 2, 3 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 1, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 1, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 2, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 64, 64, 3, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 64, 64, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 96, 96, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 96, 96, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 96, 96, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 192, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 192, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 192, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M5_PRO, GGML_TYPE_Q8_0, 576, 512, -1, 1 }, { 1, 4 } }, + + { { GGML_METAL_DEVICE_M5_MAX, GGML_TYPE_F16, 32, 32, -1, 1 }, { 2, 4 } }, { { GGML_METAL_DEVICE_M5_MAX, GGML_TYPE_F16, 32, 32, 1, 2 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M5_MAX, GGML_TYPE_F16, 32, 32, 1, 4 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M5_MAX, GGML_TYPE_F16, 32, 32, 2, 2 }, { 4, 4 } }, diff --git a/ggml/src/ggml-metal/kernels/conv.metal b/ggml/src/ggml-metal/kernels/conv.metal index 5685b5cd491..a5d5aa9d929 100644 --- a/ggml/src/ggml-metal/kernels/conv.metal +++ b/ggml/src/ggml-metal/kernels/conv.metal @@ -366,7 +366,8 @@ kernel void kernel_conv_transpose_2d( const int64_t out_x = tgpig[0]; const int64_t out_y = tgpig[1]; - const int64_t out_c = tgpig[2]; + const int64_t batch = tgpig[2] / args.OC; + const int64_t out_c = tgpig[2] % args.OC; const int64_t kw = tpitg[0]; const int64_t kh = tpitg[1]; @@ -390,7 +391,7 @@ kernel void kernel_conv_transpose_2d( if (in_x >= args.IW) continue; - const int64_t input_idx = (args.IW * args.IH) * in_c + (args.IW) * in_y + in_x; + const int64_t input_idx = (args.IW * args.IH) * (args.IC * batch + in_c) + (args.IW) * in_y + in_x; const int64_t kernel_idx = (args.KH * args.KW * args.OC) * in_c + (args.KH * args.KW) * out_c + (args.KW) * kh + kw; v += (float)src0[kernel_idx] * src1[input_idx]; @@ -408,7 +409,7 @@ kernel void kernel_conv_transpose_2d( total += shared_sum[i]; } - device float * dst_ptr = (device float *) (dst + out_x*args.nb0 + out_y * args.nb1 + out_c*args.nb2); + device float * dst_ptr = (device float *) (dst + batch*args.nb3 + out_c*args.nb2 + out_y * args.nb1 + out_x*args.nb0); dst_ptr[0] = total; } } diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 64f3325b2a5..426aac52316 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -903,6 +903,8 @@ struct ggml_backend_opencl_context { cl_kernel kernel_gemv_moe_mxfp4_f32_ns_wimg = nullptr; // weight-as-texture MoE decode GEMV cl_kernel kernel_gemm_moe_mxfp4_q8_1_dp4a = nullptr; // dp4a (int8) mxfp4 MoE prefill GEMM cl_kernel kernel_gemm_moe_q4_0_q8_1_dp4a = nullptr; // dp4a (int8) q4_0 MoE prefill GEMM + cl_kernel kernel_gemm_moe_mxfp4_q8_1_dp4a_bin = nullptr; // binary dp4a (int8) mxfp4 MoE prefill GEMM + cl_kernel kernel_gemm_moe_q4_0_q8_1_dp4a_bin = nullptr; // binary dp4a (int8) q4_0 MoE prefill GEMM cl_kernel kernel_moe_reorder_b; cl_kernel kernel_moe_histogram, kernel_moe_scan, kernel_moe_fill, kernel_moe_scatter; cl_kernel kernel_moe_scatter_stable = nullptr; // deterministic slot assignment @@ -4248,6 +4250,24 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) { GGML_LOG_CONT("."); } + // gemm_moe_mxfp4_q8_1_dp4a_bin (dp4a prefill GEMM) + if (backend_ctx->has_integer_dot) { + size_t bin_size = 0; + backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin = nullptr; + + if (use_adreno_bin_kernels(backend_ctx)) { + const char * kernel_bin = (const char *)backend_ctx->get_adreno_bin_kernel("gemm_moe_mxfp4_q8_1_dp4a_ila", &bin_size); + if (kernel_bin && bin_size > 0) { + cl_program prog = + build_program_from_binary(backend_ctx->context, backend_ctx->device, kernel_bin, CL_moe_compile_opts, bin_size); + + CL_CHECK((backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin = clCreateKernel(prog, "kernel_gemm_moe_mxfp4_q8_1_dp4a_ila", &err), err)); + CL_CHECK(clReleaseProgram(prog)); + GGML_LOG_CONT("."); + } + } + } + // gemm_moe_q4_0_q8_1_dp4a (dp4a prefill GEMM) if (backend_ctx->has_integer_dot) { #ifdef GGML_OPENCL_EMBED_KERNELS @@ -4265,6 +4285,24 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) { GGML_LOG_CONT("."); } + // gemm_moe_q4_0_q8_1_dp4a_bin (dp4a prefill GEMM) + if (backend_ctx->has_integer_dot) { + size_t bin_size = 0; + backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a_bin = nullptr; + + if (use_adreno_bin_kernels(backend_ctx)) { + const char * kernel_bin = (const char *)backend_ctx->get_adreno_bin_kernel("gemm_moe_q4_0_q8_1_dp4a_ila", &bin_size); + if (kernel_bin && bin_size > 0) { + cl_program prog = + build_program_from_binary(backend_ctx->context, backend_ctx->device, kernel_bin, CL_moe_compile_opts, bin_size); + + CL_CHECK((backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a_bin = clCreateKernel(prog, "kernel_gemm_moe_q4_0_q8_1_dp4a_ila", &err), err)); + CL_CHECK(clReleaseProgram(prog)); + GGML_LOG_CONT("."); + } + } + } + // gemm_moe_q8_1_dp4a (generic dp4a MoE GEMM; MOE_QT=80 -> q8_0 expert variant) if (backend_ctx->has_integer_dot) { #ifdef GGML_OPENCL_EMBED_KERNELS @@ -6021,9 +6059,13 @@ static ggml_backend_opencl_context * ggml_cl_init(ggml_backend_dev_t dev) { } #ifdef GGML_OPENCL_USE_ADRENO_KERNELS - // determine whether to use Adreno xmem GEMM - backend_ctx->adreno_xmem_gemm_enabled = getenv("GGML_OPENCL_ADRENO_XMEM_GEMM") != nullptr && - backend_ctx->gpu_family == GPU_FAMILY::ADRENO; + // Adreno xmem F16xF32 GEMM, default on adreno, opt out with GGML_OPENCL_ADRENO_XMEM_GEMM=0. + // This helps models with f16 attention weights, e.g., gpt-oss-20b-f16 + { + const char * xmem_env = getenv("GGML_OPENCL_ADRENO_XMEM_GEMM"); + backend_ctx->adreno_xmem_gemm_enabled = backend_ctx->gpu_family == GPU_FAMILY::ADRENO && + (xmem_env ? atoi(xmem_env) != 0 : true); + } #endif // determine whether to use large buffer for Adreno @@ -19496,9 +19538,18 @@ static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, co // GEMM using local memory // Current BK = 16, so ne00 % 16 == 0 + // + // Certain A7X compiler (E031.41) executes kernel_mul_mm_f32_f32_l4_lm poorly; + // matrices with ne11 <= 8 appears OK. + // Fallback to the MV style kernels for A7x and ne11 > 8. + // Override with GGML_OPENCL_A7X_F32_LM_BYPASS=0. + static const char * a7x_f32lm_env = getenv("GGML_OPENCL_A7X_F32_LM_BYPASS"); + static const bool a7x_f32lm_bypass = (a7x_f32lm_env == nullptr || a7x_f32lm_env[0] != '0'); if (src1t == GGML_TYPE_F32 && ne00 % 16 == 0 && - ne11 > 1) { + ne11 > 1 && + !(a7x_f32lm_bypass && src0t == GGML_TYPE_F32 && ne11 > 8 && + backend_ctx->adreno_gen == ADRENO_GPU_GEN::A7X)) { switch(src0t) { case GGML_TYPE_F32: { kernel = backend_ctx->kernel_mul_mm_f32_f32_l4_lm; @@ -21519,7 +21570,9 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0, // dot prod has to be available use_moe_dp4a = backend_ctx->has_integer_dot && use_moe_dp4a; // bin kernel takes precedence - use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_q4_0_f32_ns_bin == nullptr; + if (backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a_bin == nullptr) { + use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_q4_0_f32_ns_bin == nullptr; + } cl_buffer_region region; region.origin = 0; @@ -21625,6 +21678,10 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0, // dp4a GEMM cl_kernel dk = backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a; + if (backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a_bin) { + dk = backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a_bin; + } + int aidx = 0; CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q4_0->q_img)); CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q4_0->d)); @@ -23463,8 +23520,10 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0, : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E); // dot prod has to be available use_moe_dp4a = backend_ctx->has_integer_dot && use_moe_dp4a; - // bin kernel takes precedence - use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_mxfp4_f32_ns_bin == nullptr; + // bin kernel takes precedence, dp4a bin kernel has higher priority than normal bin kernel + if (backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin == nullptr) { + use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_mxfp4_f32_ns_bin == nullptr; + } cl_buffer_region region; region.origin = 0; @@ -23573,6 +23632,10 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0, // dp4a GEMM cl_kernel dk = backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a; + if (backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin) { + dk = backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin; + } + int aidx = 0; CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_mxfp4->q_img)); CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_mxfp4->e)); diff --git a/ggml/src/ggml-openvino/CMakeLists.txt b/ggml/src/ggml-openvino/CMakeLists.txt index cc089b721fc..af3e0758ca2 100644 --- a/ggml/src/ggml-openvino/CMakeLists.txt +++ b/ggml/src/ggml-openvino/CMakeLists.txt @@ -1,6 +1,8 @@ find_package(OpenVINO REQUIRED COMPONENTS Runtime Threading) find_package(OpenCL REQUIRED) +message(STATUS "Found OpenVINO: ${OpenVINO_DIR} (found version \"${OpenVINO_VERSION}\")") + file(GLOB_RECURSE GGML_HEADERS_OPENVINO "*.h" "*.hpp") file(GLOB_RECURSE GGML_SOURCES_OPENVINO "*.cpp") diff --git a/ggml/src/ggml-openvino/ggml-decoder.cpp b/ggml/src/ggml-openvino/ggml-decoder.cpp index 599f41aebbd..006e005cb7a 100644 --- a/ggml/src/ggml-openvino/ggml-decoder.cpp +++ b/ggml/src/ggml-openvino/ggml-decoder.cpp @@ -357,6 +357,18 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const { break; } case GGML_OP_VIEW: { + if (m_is_static && node->src[0] != nullptr && + (node->src[0]->op == GGML_OP_GATED_DELTA_NET || node->src[0]->op == GGML_OP_CONCAT)) { + // VIEW slicing a GATED_DELTA_NET combined [attn|state] output, or the conv_input + // CONCAT. The consuming CPY/RMS_NORM op recovers the true window at runtime via + // ssm_state_size / the fixed conv kernel width, so this VIEW must stay an identity + // pass-through of the full source here too (it already is on the dynamic path); + // otherwise the generic static-mode Slice below would bake in the *captured* + // cgraph's token count, which is wrong once the compiled static model runs with a + // different token count (prefill chunk size or 1). + op_case = 1; + break; + } if (node->src[0]->op == GGML_OP_VIEW) { auto * src = node->src[0]; if (ggml_nelements(node) != ggml_nelements(src)) { @@ -408,6 +420,23 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const { } break; } + case GGML_OP_POOL_2D: { + const ggml_op_pool pool_mode = static_cast(node->op_params[0]); + switch (pool_mode) { + case GGML_OP_POOL_MAX: { + op_case = 1; + break; + } + case GGML_OP_POOL_AVG: { + op_case = 2; + break; + } + default: + op_case = 0; + break; + } + break; + } case GGML_OP_CPY: { if (node->src[0]->op == GGML_OP_VIEW) { if (node->src[0]->src[0]->op == GGML_OP_GATED_DELTA_NET) { @@ -425,6 +454,31 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const { is_kvcache(node->src[1]->view_src, nullptr)) { // s_copy defrag remainder writeback: gathered extra state rows copied back into the cache op_case = 3; + } else if (node->src[1] != nullptr && node->src[1]->op == GGML_OP_VIEW && node->src[1]->view_src != nullptr) { + // op_case 5: KV write for decoder self-attention (dynamic write offset) + // op_case 6: KV write for encoder self-attn or cross-attn (static offset) + const ggml_tensor * kv_buf = node->src[1]->view_src; + if (kv_buf->ne[1] == 1 && kv_buf->ne[2] == 1 && kv_buf->ne[3] == 1) { + op_case = 6; + // Forward-scan the graph for a FLASH_ATTN_EXT that reads from + // the same buffer. Having a mask (src[3] != nullptr) implies + // decoder self-attention and the write offset is dynamic. + for (int i = 0; i < m_cgraph->n_nodes; i++) { + const ggml_tensor * n = m_cgraph->nodes[i]; + if (n->op != GGML_OP_FLASH_ATTN_EXT) { + continue; + } + // K (src[1]) and V (src[2]) are 3-D views whose view_src is + // the flat KV buffer we are writing to. + if ((n->src[1] != nullptr && n->src[1]->view_src == kv_buf) || + (n->src[2] != nullptr && n->src[2]->view_src == kv_buf)) { + if (n->src[3] != nullptr) { + op_case = 5; // decoder self-attention: mask present + } + break; + } + } + } } break; } @@ -448,6 +502,15 @@ int GgmlOvDecoder::compute_op_case(const ggml_tensor * node) const { } break; } + case GGML_OP_FLASH_ATTN_EXT: { + if (node->src[1] != nullptr && node->src[1]->op == GGML_OP_VIEW && node->src[1]->view_src != nullptr) { + const ggml_tensor * kv_buf = node->src[1]->view_src; + if (kv_buf->ne[1] == 1 && kv_buf->ne[2] == 1 && kv_buf->ne[3] == 1) { + op_case = (node->src[3] != nullptr) ? 1 : 2; + } + } + break; + } default: break; } @@ -479,23 +542,35 @@ std::pair GgmlOvDecoder::compute_llm_params(ggml_cgr switch (node->op) { case GGML_OP_FLASH_ATTN_EXT: - if (node->src[0] == nullptr || node->src[1] == nullptr || node->src[3] == nullptr) { + if (node->src[0] == nullptr || node->src[1] == nullptr) { return -1; } switch (node->src[1]->op) { case GGML_OP_PERMUTE: - // case 0: node op is FLASH_ATTN_EXT, src 1 not null & op is PERMUTE & the permuted tensor src is the view of cache k - if (node->src[1]->src[0] != nullptr && node->src[1]->src[0]->op == GGML_OP_VIEW) { + // case 0: src[1] is PERMUTE of a cache VIEW, mask required + if (node->src[3] != nullptr && node->src[1]->src[0] != nullptr && + node->src[1]->src[0]->op == GGML_OP_VIEW) { return 0; } break; case GGML_OP_CPY: - // case 1: node op is FLASH_ATTN_EXT, src 1 not null & op is CPY & the copied tensor src is PERMUTE & the permuted tensor src is the view of cache k - if (node->src[1]->src[0] != nullptr && node->src[1]->src[0]->op == GGML_OP_PERMUTE && - node->src[1]->src[0]->src[0] != nullptr && node->src[1]->src[0]->src[0]->op == GGML_OP_VIEW) { + // case 1: src[1] is CPY of a PERMUTE(VIEW), mask required + if (node->src[3] != nullptr && node->src[1]->src[0] != nullptr && + node->src[1]->src[0]->op == GGML_OP_PERMUTE && node->src[1]->src[0]->src[0] != nullptr && + node->src[1]->src[0]->src[0]->op == GGML_OP_VIEW) { return 1; } break; + case GGML_OP_VIEW: + // cases 4/5/6: whisper - K is a direct non-contiguous VIEW_3D of a KV cache + if (node->src[1]->view_src != nullptr) { + if (node->src[3] != nullptr) { + return 4; // decoder self-attention + } else { + return 5; // cross-attention or encoder self-attention + }; + } + break; default: break; } @@ -548,6 +623,18 @@ std::pair GgmlOvDecoder::compute_llm_params(ggml_cgr cache_k_permute = node->src[0]->src[0]->src[0]; mask = node->src[1]; break; + case 4: + case 5: { + // whisper: K is a direct VIEW_3D of the KV buffer, no PERMUTE node + auto * cache_k_view = node->src[1]; // VIEW_3D of kv_self.k or kv_cross.k` + compute_params.token_len_per_seq = node->src[0]->ne[1]; + if (attention_pattern_case == 4) { + compute_params.attention_size = cache_k_view->ne[1]; + } else { + compute_params.attention_size_static = cache_k_view->ne[1]; + } + continue; + } default: break; } @@ -654,10 +741,8 @@ std::pair GgmlOvDecoder::compute_llm_params(ggml_cgr ComputeParams::RsWriteback writeback; writeback.slot_begin = (int) (dest_view->view_offs / row_bytes); if (is_conv) { - // conv_input column the copied window starts at writeback.src_begin = (int) (node->src[0]->view_offs / node->src[0]->view_src->nb[0]); } else if (is_gdn) { - // first row of the state part of the gated-delta-net output writeback.src_begin = (int) (node->src[0]->view_offs / node->src[0]->view_src->nb[1]); } compute_params.rs_writebacks[get_tensor_ov_name(cgraph, node)] = writeback; @@ -718,11 +803,15 @@ ov::PartialShape GgmlOvDecoder::get_graph_input_shape(const ggml_tensor * op, } else if (is_kvcache(input, op)) { // kvcache input_shape = ov::PartialShape{get_shape(input)}; - if (!m_is_static) { + // Whisper.cpp uses a fixed size 1D KV buffer [N, 1, 1, 1] (GGML) or [1, 1, 1, N] (OV). + // the token fill level is handled by token_len_per_seq + dynamic mask input. + // skip dynamic dim and stateful reshape for this layout. + const bool is_flat_kv = (input->ne[1] == 1 && input->ne[2] == 1 && input->ne[3] == 1); + if (!m_is_static && !is_flat_kv) { // do not fix ctx size to make llama-bench work across test params input_shape[2] = -1; } - if (is_stateful()) { + if (is_stateful() && !is_flat_kv) { // Convert stateless KV cache layout [1, 1, seq, n_heads_kv * head_size] // to stateful layout [1, seq, n_heads_kv, head_size]. assert(input_shape.size() == 4 && input_shape[0] == 1 && input_shape[1] == 1 && @@ -738,7 +827,9 @@ ov::PartialShape GgmlOvDecoder::get_graph_input_shape(const ggml_tensor * op, input_shape = ov::PartialShape{1, 1, 1, len}; } else if (is_inp_s_copy(input, op) || is_s_copy_leaf(input)) { - input_shape = ov::PartialShape{1, 1, 1, -1}; + // On NPU the total slot count (n_seq_max) is fixed at translation time, so the s_copy + // index list has a static length; on CPU/GPU it may change across compiles (defrag). + input_shape = m_is_static ? ov::PartialShape{get_shape(input)} : ov::PartialShape{1, 1, 1, -1}; } else { input_shape = ov::PartialShape{get_shape(input)}; @@ -790,13 +881,16 @@ void GgmlOvDecoder::add_extra_inputs() { // see llama_kv_cache_unified::get_n_kv and llama_kv_cache_unified::get_padding. // 2. `n_seq_active` and `seq_active_start`, used in FLASH_ATTN_EXT to indicate the active sequences in the batch - auto create_1d_input = [this](const std::string & name, int64_t value) { - m_model_extra_inputs[name] = {ov::element::i64, ov::Shape{1}, value, !m_is_static}; + auto create_1d_input = [this](const std::string & name, int64_t value, bool force_parameter = false) { + m_model_extra_inputs[name] = {ov::element::i64, ov::Shape{1}, value, force_parameter || !m_is_static}; }; if (m_compute_params.attention_size != -1) { create_1d_input("attention_size", m_compute_params.attention_size); } + if (m_compute_params.attention_size_static != -1) { + create_1d_input("attention_size_static", m_compute_params.attention_size_static); + } if (m_compute_params.attention_size_swa != -1) { create_1d_input("attention_size_swa", m_compute_params.attention_size_swa); } @@ -809,17 +903,32 @@ void GgmlOvDecoder::add_extra_inputs() { // create_1d_input("token_len", m_compute_params.token_len_per_seq * m_compute_params.n_seq_active); if (m_compute_params.cache_rs_reset_idx != -1) { - create_1d_input("cache_rs_reset_idx", m_compute_params.cache_rs_reset_idx); - create_1d_input("cache_rs_reset_len", m_compute_params.cache_rs_reset_len); + // Whether/which cache slot to reset varies per compute call (e.g. a new sequence starting + // vs. continued decoding). can_reuse_statically() does not invalidate the cached static + // model on ComputeParams changes, so these must stay runtime Parameters even when static + // (scale.cpp op_case 1 only uses them in value comparisons, never as Slice bounds, so this + // does not reintroduce dynamic shapes). + create_1d_input("cache_rs_reset_idx", m_compute_params.cache_rs_reset_idx, /*force_parameter=*/true); + create_1d_input("cache_rs_reset_len", m_compute_params.cache_rs_reset_len, /*force_parameter=*/true); } if (m_compute_params.s_copy_active_slot_len != -1) { create_1d_input("s_copy_active_slot_len", m_compute_params.s_copy_active_slot_len); + if (m_is_static) { + // Number of real tokens in the current prefill chunk. The last chunk is padded with + // fabricated token ids; attention masks them out, but the recurrent (GDN/conv) path + // would otherwise fold them into cache_r/cache_s permanently. Varies per chunk, so it + // must stay a runtime Parameter; it is only compared against a Range or used as Gather + // indices, so it does not make any shape dynamic. + create_1d_input("chunk_valid_len", get_static_n_tokens(), /*force_parameter=*/true); + } } for (const auto & [node_name, writeback] : m_compute_params.rs_writebacks) { create_1d_input("rs_slot_begin_" + node_name, writeback.slot_begin); - create_1d_input("rs_src_begin_" + node_name, writeback.src_begin); + if (!m_is_static) { + create_1d_input("rs_src_begin_" + node_name, writeback.src_begin); + } } } @@ -1785,13 +1894,23 @@ void GgmlOvDecoder::compute_node_dynamic_dims() { auto dynamic_dim_stride = src_logical_nb[dynamic_dim_idx] / ggml_type_size(node->src[0]->type) * ggml_type_size(node->type); int matched_dim_count = 0; + int first_matched_dim = -1; for (int i = 0; i < GGML_MAX_DIMS; i++) { if (node->nb[i] == dynamic_dim_stride && node->ne[i] == node->src[0]->ne[dynamic_dim_idx]) { + if (first_matched_dim == -1) { + first_matched_dim = i; + } m_node_dynamic_dims[node] = i; matched_dim_count++; } } - if (matched_dim_count != 1) { + if (matched_dim_count > 1 && node->src[0]->ne[dynamic_dim_idx] == 1) { + // Single-token capture: every trailing dim is size 1 with the same stride, so + // the match is ambiguous. The lowest index is the real axis; the rest are + // ggml's size-1 padding. Bailing out here would bake the captured token count + // into the static prefill model, which then runs with a different one. + m_node_dynamic_dims[node] = first_matched_dim; + } else if (matched_dim_count != 1) { m_node_dynamic_dims[node] = -1; GGML_LOG_WARN("ggml-openvino: cannot determine dynamic dim for CONT node '%s', src[0]: '%s'\n", node->name, node->src[0]->name); diff --git a/ggml/src/ggml-openvino/ggml-decoder.h b/ggml/src/ggml-openvino/ggml-decoder.h index 8e39a26c8b7..74cb7385029 100644 --- a/ggml/src/ggml-openvino/ggml-decoder.h +++ b/ggml/src/ggml-openvino/ggml-decoder.h @@ -47,6 +47,7 @@ struct ComputeParams { int seq_active_start = 0; int attention_size = -1; int attention_size_swa = -1; + int attention_size_static = -1; // encoder/cross-attn KV fill level (whisper) int input_len = -1; int token_len_per_seq = -1; int past_kv_len = -1; @@ -84,14 +85,15 @@ struct ComputeParams { struct RsWriteback { int slot_begin = 0; // first cache slot written by the CPY - int src_begin = 0; // where the copied data starts in the source tensor (in rows of it) + int src_begin = 0; // first source row or column copied by the CPY }; std::map rs_writebacks; - // Offsets of the state cache writeback CPY nodes, keyed by node name. They change with the - // batch (kv head, active sequence count, token count) and, with rollback enabled - // (cparams.n_rs_seq > 0), the conv state is written back once per snapshot slot, each snapshot - // taking a different conv_input window. Passed to the cached model as runtime inputs. + // Destination slot offset of each state cache writeback CPY node, keyed by node name. It + // changes with the batch (kv head, active sequence count) and, with rollback enabled + // (cparams.n_rs_seq > 0), the conv state is written back once per snapshot slot. Passed to the + // cached model as a runtime input. Dynamic models also receive the source-side offset; static + // models use a fixed end-anchored offset in the translator. }; class GgmlOvDecoder : public ov::frontend::ggml::GgmlDecoder { diff --git a/ggml/src/ggml-openvino/ggml-openvino-extra.cpp b/ggml/src/ggml-openvino/ggml-openvino-extra.cpp index 36c749244f8..36dfa4d9471 100644 --- a/ggml/src/ggml-openvino/ggml-openvino-extra.cpp +++ b/ggml/src/ggml-openvino/ggml-openvino-extra.cpp @@ -32,6 +32,8 @@ void ggml_openvino_device_config::init() { "GGML_OPENVINO_DEVICE", "GGML_OPENVINO_CACHE_DIR", "GGML_OPENVINO_DEBUG_NODE", + "GGML_OPENVINO_COMPILED_MODEL_CACHE_DIR", + "GGML_OPENVINO_NPU_COMPILE_CONFIG", // Integer values (use ggml_openvino_getenv_int) "GGML_OPENVINO_PREFILL_CHUNK_SIZE", // Boolean toggles (treated as int flags via ggml_openvino_getenv_int) @@ -41,6 +43,9 @@ void ggml_openvino_device_config::init() { "GGML_OPENVINO_DUMP_IR", "GGML_OPENVINO_DEBUG_INPUT", "GGML_OPENVINO_DEBUG_OUTPUT", + // Force the static (NPU-shape) compute path on any device, e.g. GGML_OPENVINO_DEVICE=CPU, + // to test the static-shape translation without NPUW/real NPU hardware in the loop. + "GGML_OPENVINO_FORCE_STATIC", "GGML_OPENVINO_PRINT_CGRAPH_TENSOR_ADDRESS", "GGML_OPENVINO_ENABLE_CACHE", "GGML_OPENVINO_DISABLE_CACHE", @@ -50,7 +55,7 @@ void ggml_openvino_device_config::init() { "GGML_OPENVINO_MEMORY_OPTIMIZE", "GGML_OPENVINO_RELEASE_WEIGHTS", "GGML_OPENVINO_REDUCE_COMPILE_MEM", - "GGML_OPENVINO_COMPILED_MODEL_CACHE_DIR", + "GGML_OPENVINO_LOG_UNSUPPORTED_OPS", }; for (const char * const & env_var : env_var_names) { @@ -85,6 +90,11 @@ void ggml_openvino_device_config::init() { compile_config["NPUW_CACHE_DIR"] = cache_dir; compile_config.insert(ov::cache_mode(ov::CacheMode::OPTIMIZE_SIZE)); } + const char * compilation_mode_params = + ggml_openvino_getenv_str("GGML_OPENVINO_NPU_COMPILE_CONFIG"); + if (compilation_mode_params && strlen(compilation_mode_params) > 0) { + compile_config["NPU_COMPILATION_MODE_PARAMS"] = compilation_mode_params; + } } else if (cache_dir && strlen(cache_dir) > 0) { compile_config.insert(ov::cache_dir(cache_dir)); compile_config.insert(ov::cache_mode(ov::CacheMode::OPTIMIZE_SIZE)); diff --git a/ggml/src/ggml-openvino/ggml-openvino.cpp b/ggml/src/ggml-openvino/ggml-openvino.cpp index e299e16c778..4b1789713d1 100644 --- a/ggml/src/ggml-openvino/ggml-openvino.cpp +++ b/ggml/src/ggml-openvino/ggml-openvino.cpp @@ -908,11 +908,27 @@ static bool has_non_contiguous_view_input(const ggml_tensor * op) { } static bool is_supported_flash_attn_pattern(const ggml_tensor * op) { - // pattern of q,k,v should be q->op==PERMUTE, q->src[0]->op==VIEW, q->src[0]->src[0]->view_src==nullptr + // Each Q/K/V input must follow one of: + // PERMUTE -> VIEW -> base (view_src==nullptr) (llama KV-cache path) + // PERMUTE -> RESHAPE -> base (view_src==nullptr) (whisper Q) + // VIEW -> base (view_src==nullptr) (whisper K/V from kv_pad) for (int i = 0; i < 3; i++) { const ggml_tensor * src = op->src[i]; - if (src->op != GGML_OP_PERMUTE || src->src[0] == nullptr || src->src[0]->op != GGML_OP_VIEW || - src->src[0]->src[0] == nullptr || src->src[0]->src[0]->view_src != nullptr) { + if (src->op == GGML_OP_PERMUTE) { + if (src->src[0] == nullptr) { + return false; + } + if (src->src[0]->op != GGML_OP_VIEW && src->src[0]->op != GGML_OP_RESHAPE) { + return false; + } + if (src->src[0]->src[0] == nullptr || src->src[0]->src[0]->view_src != nullptr) { + return false; + } + } else if (src->op == GGML_OP_VIEW) { + if (src->src[0] == nullptr || src->src[0]->view_src != nullptr) { + return false; + } + } else { return false; } } @@ -1030,18 +1046,29 @@ static bool is_msa_block_mask_expansion(const ggml_tensor * op) { return tensor_name_starts_with(src, "msa_block_mask"); } -static bool is_op_unsupported_case(const ggml_tensor * op) { +namespace { +struct ggml_openvino_op_support { + bool is_supported = true; + std::string reason; + + operator bool() const { + return is_supported; + } +}; +} // namespace + +static ggml_openvino_op_support is_op_supported_case(const ggml_tensor * op) { if (is_msa_block_mask_expansion(op)) { - return true; + return {false, "MSA block mask expansion is not supported"}; } switch (op->op) { case GGML_OP_CONCAT: { if (op->type == GGML_TYPE_I64) { - return true; + return {false, "CONCAT with I64 type is not supported"}; } if (ggml_openvino_get_device_name() == "GPU" && op->type == GGML_TYPE_BF16 && has_view_op_input(op)) { - return true; + return {false, "CONCAT with BF16 type and VIEW input is not supported on GPU"}; } break; } @@ -1052,24 +1079,21 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { // OpenVINO SET translation currently supports dst layouts that match src0 strides. if (op->src[0] == nullptr || nb1 != op->src[0]->nb[1] || nb2 != op->src[0]->nb[2] || nb3 != op->src[0]->nb[3]) { - // std::cout << "Unsupported SET op with dst nb1=" << nb1 << ", nb2=" << nb2 << ", nb3=" << nb3 - // << " that does not match src0 strides nb[1]=" - // << (op->src[0] != nullptr ? std::to_string(op->src[0]->nb[1]) : "null") - // << ", nb[2]=" << (op->src[0] != nullptr ? std::to_string(op->src[0]->nb[2]) : "null") - // << ", nb[3]=" << (op->src[0] != nullptr ? std::to_string(op->src[0]->nb[3]) : "null") - // << std::endl; - return true; + return {false, "SET op with dst nb1=" + std::to_string(nb1) + ", nb2=" + std::to_string(nb2) + ", nb3=" + std::to_string(nb3) + + " that does not match src0 strides nb[1]=" + (op->src[0] != nullptr ? std::to_string(op->src[0]->nb[1]) : "null") + + ", nb[2]=" + (op->src[0] != nullptr ? std::to_string(op->src[0]->nb[2]) : "null") + + ", nb[3]=" + (op->src[0] != nullptr ? std::to_string(op->src[0]->nb[3]) : "null")}; } break; } case GGML_OP_GET_ROWS: case GGML_OP_SET_ROWS: { if (op->ne[3] != 1) { - return true; + return {false, "GET_ROWS/SET_ROWS with ne[3] != 1 (ne[3]=" + std::to_string(op->ne[3]) + ") is not supported"}; } if (op->op == GGML_OP_GET_ROWS && ggml_openvino_get_device_name() == "GPU" && op->src[0]->type == GGML_TYPE_BF16) { - return true; + return {false, "GET_ROWS with BF16 src0 is not supported on GPU"}; } if (op->ne[0] == 256 && (op->src[0]->type == GGML_TYPE_Q4_K || op->src[0]->type == GGML_TYPE_Q5_K || op->src[0]->type == GGML_TYPE_Q4_1 || op->src[0]->type == GGML_TYPE_Q5_1)) { @@ -1078,14 +1102,14 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { // make_int8_weights/make_int4_weights: dequant is done in f16, not f32, to keep the // Convert/Subtract/Multiply chain fusable into GatherMatmulCompressed/FullyConnectedCompressed // for the shared non-test code paths). - return true; + return {false, "GET_ROWS/SET_ROWS with ne[0] == 256 and type " + std::string(ggml_type_name(op->src[0]->type)) + + " rejected due to f16-arithmetic dequant rounding errors that intermittently exceed 1e-7 NMSE threshold"}; } - break; } case GGML_OP_RESHAPE: { if (strncmp(op->name, "ffn_norm_exps", sizeof("ffn_norm_exps") - 1) == 0) { - return true; + return {false, "RESHAPE for ffn_norm_exps is not supported"}; } break; } @@ -1093,11 +1117,13 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { case GGML_OP_MUL: case GGML_OP_SUB: { if (op->src[1]->op == GGML_OP_PERMUTE) { - return true; + return {false, "ADD/MUL/SUB with PERMUTE src1 is not supported"}; } for (int i = 0; i < 4; i++) { if (op->src[0]->ne[i] != op->src[1]->ne[i] && (op->src[0]->ne[i] != 1 && op->src[1]->ne[i] != 1)) { - return true; + return {false, "ADD/MUL/SUB with incompatible broadcast shapes: src0->ne[" + std::to_string(i) + "]=" + + std::to_string(op->src[0]->ne[i]) + ", src1->ne[" + std::to_string(i) + "]=" + + std::to_string(op->src[1]->ne[i])}; } } break; @@ -1106,7 +1132,7 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { // Keep support aligned with the CPU backend implementation, which only handles f32 inputs/output and i32 ids. if (op->type != GGML_TYPE_F32 || op->src[0]->type != GGML_TYPE_F32 || op->src[1]->type != GGML_TYPE_F32 || op->src[2]->type != GGML_TYPE_I32) { - return true; + return {false, "ADD_ID only supports F32 inputs/output and I32 ids"}; } break; } @@ -1116,14 +1142,27 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { // until the fused GPU kernel is reliable. (falied case llama-arch-test mpt) if (ggml_openvino_get_device_name() == "GPU" && op->src[1]->ne[0] == op->ne[0] && op->src[1]->ne[1] == 1 && op->src[1]->ne[2] == 1 && op->src[1]->ne[3] == 1) { - return true; + return {false, "DIV per-channel scale broadcast is not supported on GPU"}; + } + break; + } + case GGML_OP_POOL_2D: { + const auto& name = ggml_openvino_get_device_name(); + if (name == "GPU") { + const int32_t * params = op->op_params; + const int k0 = params[1]; + const int k1 = params[2]; + const int p0 = params[5]; + const int p1 = params[6]; + if ((p0 > 0 || p1 > 0) && (k0 < 3 || k1 < 3)) { + return {false, "POOL_2D with padding and kernel size < 3 is not supported on " + name}; + } } break; } case GGML_OP_SUM_ROWS: { - // if the input is PERMUTE skip if (op->src[0]->op == GGML_OP_PERMUTE) { - return true; + return {false, "SUM_ROWS with PERMUTE input is not supported"}; } break; } @@ -1140,54 +1179,51 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { // accuracy drift in the OpenVINO path. Restrict by scale=1.0 to avoid // affecting non-gemma3n models such as Llama-3.2. if (fabsf(scale - 1.0f) < 1e-6f && is_gemma3n_flash_attn_pattern(op)) { - return true; + return {false, "FLASH_ATTN_EXT gemma3n pattern on GPU is not supported"}; } if (op->src[4] != nullptr) { - // GGML_LOG_WARN("OpenVINO backend does not support FLASH_ATTN_EXT with sinks\n"); - return true; + return {false, "FLASH_ATTN_EXT with sinks is not supported"}; } if (!is_supported_flash_attn_pattern(op)) { - return true; + return {false, "FLASH_ATTN_EXT unsupported attention pattern"}; } if (max_bias > 0) { - // GGML_LOG_WARN("OpenVINO backend does not support FLASH_ATTN_EXT with max_bias > 0\n"); - return true; + return {false, "FLASH_ATTN_EXT with max_bias > 0 (max_bias=" + std::to_string(max_bias) + ") is not supported"}; } if (logit_softcap != 0) { - // GGML_LOG_WARN("OpenVINO backend does not support FLASH_ATTN_EXT with logit_softcap != 0\n"); - return true; + return {false, "FLASH_ATTN_EXT with logit_softcap != 0 (logit_softcap=" + std::to_string(logit_softcap) + ") is not supported"}; } break; } case GGML_OP_PERMUTE: { - if (op->type == GGML_TYPE_BF16) { - // err msg: [GPU] Could not find a suitable kernel for transpose - // GGML_LOG_WARN("OpenVINO backend does not support PERMUTE with BF16 type\n"); - return true; + if (op->type == GGML_TYPE_BF16 && ggml_openvino_get_device_name() == "GPU") { + return {false, "PERMUTE with BF16 type is not supported on GPU"}; } break; } case GGML_OP_CPY: { if (op->src[0]->type == GGML_TYPE_BF16 || op->src[1]->type == GGML_TYPE_BF16) { - // GGML_LOG_WARN("OpenVINO backend does not support CPY with non-contiguous data or bf16 types\n"); - return true; + return {false, "CPY with BF16 src type is not supported"}; } // CPY to a quantized destination (e.g. f32 -> q4_0) is numerically unstable with OpenVINO backend. if (ggml_is_quantized(op->type)) { - return true; + return {false, "CPY to quantized destination (e.g. f32 -> q4_0) is numerically unstable"}; } if (ggml_nelements(op->src[0]) != ggml_nelements(op->src[1])) { - return true; + return {false, "CPY with mismatched element counts is not supported: src0=" + std::to_string(ggml_nelements(op->src[0])) + + " != src1=" + std::to_string(ggml_nelements(op->src[1]))}; } // op test case with non-contiguous src or dst if ((op->ne[0] == 3 && op->ne[1] == 4 && op->ne[2] == 3 && op->ne[3] == 2) || (op->ne[0] == 1 && op->ne[1] == 4 && op->ne[2] == 3 && op->ne[3] == 2) || (op->ne[0] == 2 && op->ne[1] == 4 && op->ne[2] == 3 && op->ne[3] == 2)) { - return true; + return {false, "CPY with non-contiguous shape [" + std::to_string(op->ne[0]) + ", " + + std::to_string(op->ne[1]) + ", " + std::to_string(op->ne[2]) + ", " + + std::to_string(op->ne[3]) + "] is not supported"}; } if (!cpy_output_view_is_supported(op)) { - return true; + return {false, "CPY with non-contiguous output view is not supported"}; } break; } @@ -1196,13 +1232,14 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { ggml_is_quantized(op->src[0]->type) && strcmp(op->src[0]->name, "a") == 0 && strcmp(op->src[1]->name, "b") == 0 && op->src[0]->ne[1] == 1 && op->src[1]->ne[1] == 64 && op->src[0]->ne[0] == 256 && op->src[1]->ne[0] == 256) { - return true; + return {false, "MUL_MAT quantized benchmark test case on GPU is not supported"}; } if (op->src[0]->ne[3] != op->src[1]->ne[3] && op->src[0]->ne[3] != 1 && op->src[1]->ne[3] != 1) { - return true; + return {false, "MUL_MAT with incompatible broadcast on ne[3]: src0->ne[3]=" + std::to_string(op->src[0]->ne[3]) + + ", src1->ne[3]=" + std::to_string(op->src[1]->ne[3])}; } if (op->src[0]->op == GGML_OP_VIEW && op->src[1]->op == GGML_OP_VIEW) { - return true; + return {false, "MUL_MAT with both inputs as VIEW is not supported"}; } break; } @@ -1210,16 +1247,17 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { // Single-expert (or empty) MUL_MAT_ID is a degenerate shape that stresses GatherMatmul edge // cases and never occurs in real MoE; let it fall back to CPU. if (op->src[0] != nullptr && op->src[0]->ne[2] <= 1) { - return true; + return {false, "MUL_MAT_ID with single-expert or empty ne[2] <= 1 (ne[2]=" + + std::to_string(op->src[0]->ne[2]) + ") is not supported"}; } if (ggml_openvino_get_device_name() == "GPU" && op->src[0] != nullptr && op->src[0]->type == GGML_TYPE_BF16) { - return true; + return {false, "MUL_MAT_ID with BF16 weights on GPU is not supported"}; } // GPU MUL_MAT_ID uses a Gather+MatMul fallback because the GPU plugin rejects internal // GatherMatmul for these test shapes. Skip cases that would materialize a large selected // expert-weight temporary. if (ggml_openvino_get_device_name() == "GPU" && mul_mat_id_requires_large_tmp(op)) { - return true; + return {false, "MUL_MAT_ID requires large temporary on GPU"}; } break; } @@ -1229,51 +1267,46 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { const int mode = op_params[2]; if (op_params[15] != 0) { // FIXME: support ggml_rope_set_offset - return true; + return {false, "ggml_rope_set_offset is not supported"}; } if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX && mode != GGML_ROPE_TYPE_IMROPE) { - // GGML_LOG_WARN("OpenVINO backend does not support ROPE with mode %d\n", mode); - return true; + return {false, "ROPE with mode " + std::to_string(mode) + " is not supported"}; } const int64_t head_dim = op->src[0]->ne[0]; const int64_t rope_dims = n_dims == 0 ? head_dim : n_dims; if (rope_dims <= 0 || rope_dims > head_dim || (rope_dims % 2) != 0) { - // GGML_LOG_WARN("OpenVINO backend does not support ROPE with n_dims %d and src[0]->ne[0] %ld\n", n_dims, - // op->src[0]->ne[0]); - return true; + return {false, "ROPE with n_dims=" + std::to_string(n_dims) + ", head_dim=" + std::to_string(head_dim) + " is not supported"}; } if (op->type != GGML_TYPE_F32 && op->type != GGML_TYPE_F16) { - // GGML_LOG_WARN("OpenVINO backend does not support ROPE with type %s\n", ggml_type_name(op->type)); - return true; + return {false, "ROPE with type " + std::string(ggml_type_name(op->type)) + " is not supported"}; } if (op->src[0]->op == GGML_OP_VIEW) { - if (op->src[0]->view_src->ne[1] != op->src[0]->ne[2]) { - // GGML_LOG_WARN( - // "OpenVINO backend does not support ROPE with src[0]->view_src->ne[1] %ld != src[0]->ne[2] " - // "%ld\n", - // op->src[0]->view_src->ne[1], op->src[0]->ne[2]); - return true; + const struct ggml_tensor * view = op->src[0]; + const struct ggml_tensor * view_src = view->view_src; + if (view_src->ne[1] != view->ne[1] || view_src->ne[2] != view->ne[2] || view_src->ne[3] != view->ne[3]) { + return {false, "ROPE with view_src->ne [" + std::to_string(view_src->ne[1]) + ", " + + std::to_string(view_src->ne[2]) + ", " + std::to_string(view_src->ne[3]) + + "] != view->ne [" + std::to_string(view->ne[1]) + ", " + + std::to_string(view->ne[2]) + ", " + std::to_string(view->ne[3]) + + "] is not supported"}; } } if (mode == GGML_ROPE_TYPE_IMROPE && (op->src[2] != 0 || ((const float *) op_params)[6] != 1 || ((const float *) op_params)[7] != 0 || ((const float *) op_params)[8] != 1)) { - // GGML_LOG_WARN("OpenVINO backend does not support IMROPE with freq_factors, freq_scale, ext_factor, and attn_factor\n"); - return true; + return {false, "IMROPE with freq_factors, freq_scale, ext_factor, and attn_factor is not supported"}; } break; } case GGML_OP_TRANSPOSE: { - // if the type is bf16, will return true if (op->type == GGML_TYPE_BF16) { - // GGML_LOG_WARN("OpenVINO backend does not support CONT with BF16 type\n"); - return true; + return {false, "TRANSPOSE with BF16 type is not supported"}; } break; } case GGML_OP_REPEAT: { if (ggml_openvino_get_device_name() == "GPU" && op->type == GGML_TYPE_BF16) { - return true; + return {false, "REPEAT with BF16 type is not supported on GPU"}; } break; } @@ -1285,15 +1318,15 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { // return true; // } if (op->src[2]->op == GGML_OP_PERMUTE) { - return true; + return {false, "GATED_DELTA_NET with PERMUTE src2 is not supported"}; } // kda (per-key-dimension gating) not supported by fused GatedDeltaNet op if (op->src[3]->ne[0] != 1) { - return true; + return {false, "GATED_DELTA_NET with kda (per-key-dimension gating) is not supported"}; } // K > 1 (multiple state snapshots) not supported by fused op if (((const int32_t *) op->op_params)[0] > 1) { - return true; + return {false, "GATED_DELTA_NET with K > 1 (multiple state snapshots) is not supported"}; } break; } @@ -1307,17 +1340,17 @@ static bool is_op_unsupported_case(const ggml_tensor * op) { // Skip TOPK_MOE fused tests until it is fully supported. // The argsort_top_k VIEW wrapping ARGSORT is named "selected_experts" in test_topk_moe. if (strcmp(op->name, "selected_experts") == 0) { - return true; + return {false, "VIEW for selected_experts (argsort_top_k) is not supported"}; } break; } default: break; } - return false; + return {true, ""}; } -static bool ggml_backend_openvino_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { +static ggml_openvino_op_support ggml_backend_openvino_device_supports_op_impl(ggml_backend_dev_t dev, const ggml_tensor * op) { GGML_ASSERT(dev->reg != nullptr); static std::unordered_set supported_types{ @@ -1367,48 +1400,41 @@ static bool ggml_backend_openvino_device_supports_op(ggml_backend_dev_t dev, con case GGML_OP_UNARY: { auto supported = supported_unary_ops.find(ggml_get_unary_op(op)) != supported_unary_ops.end(); if (!supported) { - // GGML_LOG_WARN("OpenVINO backend does not support unary op %s\n", ggml_unary_op_name(ggml_get_unary_op(op))); - return false; + return {false, "unary op " + std::string(ggml_unary_op_name(ggml_get_unary_op(op))) + " has no op translator"}; } if (ggml_get_unary_op(op) == GGML_UNARY_OP_EXP && op->type == GGML_TYPE_F32) { - return false; + return {false, "UNARY_EXP with F32 type is not supported"}; } break; } case GGML_OP_GLU: { auto supported = supported_glu_ops.find(ggml_get_glu_op(op)) != supported_glu_ops.end(); if (!supported) { - // GGML_LOG_WARN("OpenVINO backend does not support GLU op %s\n", ggml_glu_op_name(ggml_get_glu_op(op))); - return false; + return {false, "GLU op " + std::string(ggml_glu_op_name(ggml_get_glu_op(op))) + " has no op translator"}; } // if (has_view_op_input(op)) { - // // GGML_LOG_WARN("OpenVINO backend does not support unary op %s with view input\n", - // // ggml_glu_op_name(ggml_get_glu_op(op))); - // return false; + // return {false, "GLU op " + std::string(ggml_glu_op_name(ggml_get_glu_op(op))) + " with view input is not supported"}; // } if (op->src[1] == nullptr && op->src[0]->ne[0] % 2 != 0) { // triggers bug in ov gpu - return false; + return {false, "GLU op with odd src0 ne[0] and null src1 is not supported"}; } break; } default: { auto supported = supported_ops.find(op->op) != supported_ops.end(); if (!supported) { - // GGML_LOG_WARN("OpenVINO backend does not support op %s\n", ggml_op_name(op->op)); - return false; + return {false, "op " + std::string(ggml_op_name(op->op)) + " has no op translator"}; } static std::set ops_not_support_view_input{}; if (ops_not_support_view_input.find(op->op) != ops_not_support_view_input.end() && has_view_op_input(op)) { - // GGML_LOG_WARN("OpenVINO backend does not support op %s with view input\n", ggml_op_name(op->op)); - return false; + return {false, "op " + std::string(ggml_op_name(op->op)) + " with VIEW input is not supported"}; } } } if (supported_types.find(op->type) == supported_types.end()) { - // GGML_LOG_WARN("OpenVINO backend does not support tensor type %s\n", ggml_type_name(op->type)); - return false; + return {false, "tensor type " + std::string(ggml_type_name(op->type)) + " is not supported"}; } for (int i = 0; i < GGML_MAX_SRC; i++) { auto * src = op->src[i]; @@ -1416,21 +1442,32 @@ static bool ggml_backend_openvino_device_supports_op(ggml_backend_dev_t dev, con break; } if (supported_types.find(src->type) == supported_types.end()) { - // GGML_LOG_WARN("OpenVINO backend does not support tensor type %s\n", ggml_type_name(src->type)); - return false; + return {false, "src[" + std::to_string(i) + "] type " + std::string(ggml_type_name(src->type)) + " is not supported"}; } const bool is_supported_3d_moe_expert = op->op == GGML_OP_MUL_MAT_ID && i == 0 && (src->type == GGML_TYPE_MXFP4 || src->ne[3] == 1); if (ggml_is_quantized(src->type) && src->ne[2] != 1 && !is_supported_3d_moe_expert) { - // GGML_LOG_WARN("OpenVINO backend does not support 3D quantized tensors\n"); - return false; + return {false, "3D quantized tensor for src[" + std::to_string(i) + "] is not supported"}; } } - if (is_op_unsupported_case(op)) { - return false; + auto op_support_case = is_op_supported_case(op); + if (!op_support_case.is_supported) { + return op_support_case; } - return true; + return {true, ""}; +} + +static bool ggml_backend_openvino_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { + auto res = ggml_backend_openvino_device_supports_op_impl(dev, op); + if (!res.is_supported) { + static const bool log_unsupported = ggml_openvino_getenv_int("GGML_OPENVINO_LOG_UNSUPPORTED_OPS") != 0; + if (log_unsupported) { + GGML_LOG_WARN("OpenVINO op unsupported: op '%s' (%s), type %s: %s\n", + op->name, ggml_op_name(op->op), ggml_type_name(op->type), res.reason.c_str()); + } + } + return res.is_supported; } static bool ggml_backend_openvino_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { diff --git a/ggml/src/ggml-openvino/openvino/op/cpy.cpp b/ggml/src/ggml-openvino/openvino/op/cpy.cpp index 5b387fc50d3..6f1e34779ac 100644 --- a/ggml/src/ggml-openvino/openvino/op/cpy.cpp +++ b/ggml/src/ggml-openvino/openvino/op/cpy.cpp @@ -3,8 +3,11 @@ #include "../utils.h" #include +#include +#include #include -#include +#include +#include #include #include #include @@ -12,9 +15,14 @@ #include #include #include +#include #include +#include #include #include +#include +#include +#include namespace ov { namespace frontend { @@ -61,10 +69,27 @@ OutputVector translate_cpy(const NodeContext & context) { return rename_outputs_with_suffix({res}, context.get_name()); } - // Recurrent state cache writeback into a slot block of the cache. Where the block starts and - // where the copied data starts in the source are runtime inputs, so the cached model works for - // any kv head, active sequence count and token count. The result is the full updated cache. + // Recurrent state cache writeback into a slot block of the cache. Where the block starts is a + // runtime input, so the cached model works for any kv head and active sequence count. The + // result is the full updated cache. // op_case 1: gated-delta-net state, op_case 2: conv state, op_case 3: defrag remainder. + if (op_case == 3) { + // With -np 1 (and generally whenever there is no defrag remainder) this GET_ROWS gathers + // zero rows: nothing to write back, and the cache is unchanged. NPU rejects zero-size + // tensors, so short-circuit instead of building a degenerate Slice/Concat chain. + bool is_empty = false; + if (input_shape.rank().is_static()) { + for (const auto & d : input_shape) { + if (d.is_static() && d.get_length() == 0) { + is_empty = true; + break; + } + } + } + if (is_empty) { + return {context.get_input(1)}; + } + } const std::string slot_begin_name = "rs_slot_begin_" + context.get_name(); const bool slice_assign = context.has_input(slot_begin_name) && !context.is_stateful() && (op_case >= 1 && op_case <= 3); @@ -81,19 +106,49 @@ OutputVector translate_cpy(const NodeContext & context) { ov::Output begin = context.get_input(slot_begin_name); auto base = context.get_input(1); if (op_case == 1) { - // GDN packs [attn | state snapshots]; the state part runs from src_begin to the end. - auto src_begin = context.get_input("rs_src_begin_" + context.get_name()); - auto state_part = std::make_shared(context.get_input(0), src_begin, int_max, one, axis); + ov::Output state_begin; + const std::string src_begin_name = "rs_src_begin_" + context.get_name(); + if (context.has_input(src_begin_name)) { + state_begin = context.get_input(src_begin_name); + } else { + auto ssm_state_size = context.get_ssm_state_size(); + if (context.has_input("s_copy_active_slot_len")) { + auto len = context.get_input("s_copy_active_slot_len"); + auto state_rows = std::make_shared( + ov::op::v0::Constant::create(ov::element::i64, {1}, {ssm_state_size}), len); + state_begin = std::make_shared(state_rows); + } else { + state_begin = ov::op::v0::Constant::create(ov::element::i64, {1}, {-ssm_state_size}); + } + } + auto state_part = + std::make_shared(context.get_input(0), state_begin, int_max, one, axis); src = std::make_shared(state_part, feature, false); } else if (op_case == 2) { - // conv_input is [previous conv state | new tokens]; copy the conv_kernel_size - 1 wide - // window starting at src_begin, which is the snapshot this writeback corresponds to. + // conv_input is [previous conv state | new tokens]; the snapshot is the conv_kernel_size - 1 + // columns ending at the last *valid* token. Gather (rather than Slice) keeps the output + // shape static even though the window start is a runtime value. auto window_size = (int64_t) input_shape[3].get_length(); - auto src_begin = context.get_input("rs_src_begin_" + context.get_name()); - auto src_end = std::make_shared( - src_begin, ov::op::v0::Constant::create(ov::element::i64, {1}, {window_size})); - auto window = std::make_shared(context.get_input(0), src_begin, src_end, one, - ov::op::v0::Constant::create(ov::element::i64, {1}, {3})); + ov::Output window; + auto col_axis = ov::op::v0::Constant::create(ov::element::i64, {1}, {3}); + const std::string src_begin_name = "rs_src_begin_" + context.get_name(); + if (context.has_input(src_begin_name)) { + auto src_begin = context.get_input(src_begin_name); + auto src_end = std::make_shared( + src_begin, ov::op::v0::Constant::create(ov::element::i64, {1}, {window_size})); + window = std::make_shared(context.get_input(0), src_begin, src_end, one, col_axis); + } else if (context.has_input("chunk_valid_len")) { + std::vector offsets(window_size); + std::iota(offsets.begin(), offsets.end(), 0); + auto indices = std::make_shared( + ov::op::v0::Constant::create(ov::element::i64, {(size_t) window_size}, offsets), + context.get_input("chunk_valid_len")); + window = std::make_shared(context.get_input(0), indices, col_axis); + } else { + auto window_begin = ov::op::v0::Constant::create(ov::element::i64, {1}, {-window_size}); + window = + std::make_shared(context.get_input(0), window_begin, int_max, one, col_axis); + } const auto base_shape = base.get_partial_shape(); FRONT_END_OP_CONVERSION_CHECK(base_shape.rank().is_static() && base_shape.rank().get_length() == 4, "CPY conv state cache update requires rank-4 base cache"); @@ -157,6 +212,63 @@ OutputVector translate_cpy(const NodeContext & context) { auto input = process_view_input_new(context, 0); + if (op_case == 5 || op_case == 6) { + auto input_shape = context.get_input_shape(0); + auto output_shape = context.get_output_shape(); + auto dst_ggml_shape = context.get_view_input_ggml_shape(1, 0); + auto dst_stride = context.get_view_input_stride(1, 0); + size_t offset_bytes = context.get_view_input_offset(1, 0); + auto n_state = (int64_t) context.get_input_shape(0)[3].get_length(); + auto n_state_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {n_state}); + auto kv_buf = context.get_input(1); // shape {1,1,1,N} + + Output token_len_per_seq; + Output n_write_dyn; + if (context.has_input("token_len_per_seq")) { + token_len_per_seq = context.get_input("token_len_per_seq"); + n_write_dyn = std::make_shared(token_len_per_seq, n_state_c); + } else { + n_write_dyn = ov::op::v0::Constant::create(ov::element::i64, {1}, {(int64_t) dst_ggml_shape[3]}); + } + size_t elem_size = dst_stride[3]; + FRONT_END_OP_CONVERSION_CHECK(elem_size > 0, "CPY KV cache view update has invalid element size"); + int64_t start_elem = (int64_t) (offset_bytes / elem_size); + // op_case 5: decoder self-attention – write offset advances each step. + // op_case 6: encoder self-attn or cross-attn – offset fixed at compile time. + const bool is_decoder_self_attn = (op_case == 5); + auto ones_c = ov::op::v0::Constant::create(ov::element::i64, {3}, std::vector{1, 1, 1}); + auto new_shape = std::make_shared(ov::OutputVector{ones_c, n_write_dyn}, 0); + + auto reshaped = std::make_shared(input, new_shape, false); + auto data = std::make_shared(reshaped, context.get_output_type()); + // Indices [start_elem .. start_elem + n_write) on axis 3 of {1,1,1,N} + // For decoder self-attention the write offset advances each step, so compute it + // dynamically from the model inputs: start = (attention_size - token_len_per_seq) * n_state. + // For encoder self-attn and cross-attn the offset is fixed at graph-compile time. + ov::Output start; + if (is_decoder_self_attn && context.has_input("attention_size") && context.has_input("token_len_per_seq")) { + auto attention_size_in = context.get_input("attention_size"); + auto token_len_in = context.get_input("token_len_per_seq"); + auto past_tokens = std::make_shared(attention_size_in, token_len_in); + auto new_start = std::make_shared(past_tokens, n_state_c); + start = std::make_shared( + new_start, ov::op::v0::Constant::create(ov::element::i64, {1}, {start_elem})); + } else { + start = ov::op::v0::Constant::create(ov::element::i64, {1}, {start_elem}); + } + auto start_squeezed = std::make_shared(start); + auto end = std::make_shared(start_squeezed, n_write_dyn); + auto end_squeezed = std::make_shared(end); + auto step = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); + auto step_squeezed = std::make_shared(step); + auto indices = + std::make_shared(start_squeezed, end_squeezed, step_squeezed, ov::element::i64); + auto axis = ov::op::v0::Constant::create(ov::element::i64, {1}, {3}); + + auto kv_updated = std::make_shared(kv_buf, indices, data, axis); + return rename_outputs_with_suffix({kv_updated}, context.get_name()); + } + if (input_shape != output_shape) { auto new_shape = ov::op::v0::Constant::create( ov::element::i64, {static_cast(output_shape.rank().get_length())}, output_shape.to_shape()); diff --git a/ggml/src/ggml-openvino/openvino/op/flash_attn_ext.cpp b/ggml/src/ggml-openvino/openvino/op/flash_attn_ext.cpp index 582df0130b5..06547f3d296 100644 --- a/ggml/src/ggml-openvino/openvino/op/flash_attn_ext.cpp +++ b/ggml/src/ggml-openvino/openvino/op/flash_attn_ext.cpp @@ -3,8 +3,8 @@ #include "../utils.h" #include "ggml-openvino/ggml-openvino-extra.h" +#include #include -#include #include #include #include @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -24,13 +25,62 @@ namespace ov { namespace frontend { namespace ggml { namespace op { +static ov::Output reshape_flat_kv(const ov::Output & kv_flat, + size_t view_offset_bytes, + size_t nb1_bytes, + int64_t n_head, + int64_t head_size, + const ov::Output & attention_size) { + int64_t n_state = n_head * head_size; + int64_t layer_start_elem = (int64_t) (view_offset_bytes / (nb1_bytes / n_state)); + // Dynamic slice: [layer_start_elem, layer_start_elem + n_kv * n_state) + auto start_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {layer_start_elem}); + auto n_state_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {n_state}); + // end = start + attention_size * n_state (both static + dynamic) + auto kv_len_elems = std::make_shared(attention_size, n_state_c); + auto end_c = std::make_shared(start_c, kv_len_elems); + auto step_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); + auto axis_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {3}); + auto sliced = std::make_shared(kv_flat, start_c, end_c, step_c, axis_c); + + // KV cache is laid out as {n_kv, n_head, head_size} in memory + // Reshape to {1, n_kv, n_head, head_size}, then transpose to {1, n_head, n_kv, head_size} + // as required by SDPA. + auto one_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); + auto n_head_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {n_head}); + auto head_size_c = ov::op::v0::Constant::create(ov::element::i64, {1}, {head_size}); + // reshape: {n_kv*n_state} -> {1, n_kv, n_head, head_size} + auto new_shape = + std::make_shared(ov::OutputVector{one_c, attention_size, n_head_c, head_size_c}, 0); + auto reshaped = std::make_shared(sliced, new_shape, false); + // transpose: {1, n_kv, n_head, head_size} -> {1, n_head, n_kv, head_size} + auto perm = ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 2, 1, 3}); + auto ret = std::make_shared(reshaped, perm); + return ret; +} OutputVector translate_flash_attn_ext(const NodeContext & context) { - num_inputs_check(context, 4, 4); + num_inputs_check(context, 3, 4); + const bool has_mask = context.get_input_size() == 4; auto q_f32 = context.get_input(0); auto k = context.get_input(1); auto v = context.get_input(2); - auto mask = context.get_input(3); + const int op_case = context.get_op_case(); + + if (op_case == 1 || op_case == 2) { + int64_t n_state_head = (int64_t) context.get_view_input_ggml_shape(1, 0)[3]; + int64_t n_head = (int64_t) context.get_view_input_ggml_shape(1, 0)[1]; + size_t nb1 = context.get_view_input_stride(1, 0)[2]; + size_t offset = context.get_view_input_offset(1, 0); + ov::Output attention_size; + if (op_case == 1) { + attention_size = context.get_input("attention_size"); + } else { + attention_size = context.get_input("attention_size_static"); + } + k = reshape_flat_kv(k, offset, nb1, n_head, n_state_head, attention_size); + v = reshape_flat_kv(v, offset, nb1, n_head, n_state_head, attention_size); + } float * params = reinterpret_cast(context.get_output_op_params()); float scale = params[0]; @@ -43,16 +93,19 @@ OutputVector translate_flash_attn_ext(const NodeContext & context) { ov::Output res; // For stateful - std::string mask_name = "KQ_mask_sliced"; - if (context.get_input_names()[3].find("swa") != std::string::npos) { - mask_name = "KQ_mask_swa_sliced"; - } - if (context.has_input(mask_name)) { - mask = context.get_input(mask_name); - } - - if (mask.get_element_type() != ov::element::f16) { - mask = std::make_shared(mask, ov::element::f16); + ov::Output mask; + if (has_mask) { + mask = context.get_input(3); + std::string mask_name = "KQ_mask_sliced"; + if (context.get_input_names()[3].find("swa") != std::string::npos) { + mask_name = "KQ_mask_swa_sliced"; + } + if (context.has_input(mask_name)) { + mask = context.get_input(mask_name); + } + if (mask.get_element_type() != ov::element::f16) { + mask = std::make_shared(mask, ov::element::f16); + } } //auto tile_kv = [&](int64_t num_heads, int64_t num_heads_kv, int64_t head_size, ov::Output kv) { @@ -108,10 +161,14 @@ OutputVector translate_flash_attn_ext(const NodeContext & context) { // get [B, 1, 1, S_q, S_k], which NUMPY-broadcasts cleanly against the // [B, num_heads_kv, factor, S_q, S_k] scores: B==B, then 1→num_heads_kv and // 1→factor on the head dims. - auto mask_unsq1 = - std::make_shared(mask, ov::op::v0::Constant::create(ov::element::i64, {1}, {2})); - // mask_unsq1: [B, 1, 1, S_q, S_k] (rank 5) - ov::Output qk_masked = std::make_shared(qk_scaled, mask_unsq1); + ov::Output qk_masked; + if (has_mask) { + auto mask_unsq1 = + std::make_shared(mask, ov::op::v0::Constant::create(ov::element::i64, {1}, {2})); + qk_masked = std::make_shared(qk_scaled, mask_unsq1); + } else { + qk_masked = qk_scaled; + } auto softmax = std::make_shared(qk_masked, /*axis=*/-1); @@ -164,9 +221,16 @@ OutputVector translate_flash_attn_ext(const NodeContext & context) { k = tile_kv(num_heads, num_heads_kv, head_size, k); v = tile_kv(num_heads, num_heads_kv, head_size, v); - auto sdpa = std::make_shared(q, k, v, mask, scale_node, false); - res = std::make_shared(sdpa, - ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 2, 1, 3})); + constexpr auto causal = false; + if (has_mask) { + auto sdpa = std::make_shared(q, k, v, mask, scale_node, causal); + res = std::make_shared( + sdpa, ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 2, 1, 3})); + } else { + auto sdpa = std::make_shared(q, k, v, scale_node, causal); + res = std::make_shared( + sdpa, ov::op::v0::Constant::create(ov::element::i64, {4}, {0, 2, 1, 3})); + } res = std::make_shared(res, ov::element::f32); return rename_outputs_with_suffix({res}, context.get_name()); } diff --git a/ggml/src/ggml-openvino/openvino/op/gated_delta_net.cpp b/ggml/src/ggml-openvino/openvino/op/gated_delta_net.cpp index 66c74828331..07eeb3c8fd6 100644 --- a/ggml/src/ggml-openvino/openvino/op/gated_delta_net.cpp +++ b/ggml/src/ggml-openvino/openvino/op/gated_delta_net.cpp @@ -7,12 +7,15 @@ #include #include #include +#include #include #include #include #include +#include #include #include +#include #include #include #include @@ -80,6 +83,28 @@ OutputVector translate_gated_delta_net(const NodeContext & context) { g = std::make_shared(g, ov::op::v0::Constant::create(ov::element::i64, {1}, {3})); beta = std::make_shared(beta, ov::op::v0::Constant::create(ov::element::i64, {1}, {3})); + if (context.has_input("chunk_valid_len")) { + // The last prefill chunk is padded with fabricated tokens. The recurrence is + // S_t = S_{t-1} * exp(g_t) + k_t (x) ((v_t - S_{t-1}^T k_t) * beta_t) + // so forcing g = 0 and beta = 0 makes a padded step an exact identity and keeps the final + // state equal to the state after the last real token. Attention output at those positions + // is garbage but never read. + const auto & g_ps = g.get_partial_shape(); + FRONT_END_OP_CONVERSION_CHECK(g_ps.rank().is_static() && g_ps.rank().get_length() == 3 && g_ps[1].is_static(), + "GATED_DELTA_NET pad masking requires a static token dimension"); + const int64_t n_tokens = g_ps[1].get_length(); + std::vector positions(n_tokens); + std::iota(positions.begin(), positions.end(), 0); + auto valid = std::make_shared( + ov::op::v0::Constant::create(ov::element::i64, {(size_t) n_tokens}, positions), + context.get_input("chunk_valid_len")); + auto mask = std::make_shared( + std::make_shared(valid, g.get_element_type()), + ov::op::v0::Constant::create(ov::element::i64, {2}, std::vector{0, 2})); + g = std::make_shared(g, mask); + beta = std::make_shared(beta, mask); + } + // std::cout << "GatedDeltaNet input shapes: q=" << q.get_partial_shape() << ", k=" << k.get_partial_shape() // << ", v=" << v.get_partial_shape() << ", g=" << g.get_partial_shape() // << ", beta=" << beta.get_partial_shape() << ", state=" << state.get_partial_shape() << std::endl; diff --git a/ggml/src/ggml-openvino/openvino/op/glu_geglu_quick.cpp b/ggml/src/ggml-openvino/openvino/op/glu_geglu_quick.cpp new file mode 100644 index 00000000000..c6d64aed43a --- /dev/null +++ b/ggml/src/ggml-openvino/openvino/op/glu_geglu_quick.cpp @@ -0,0 +1,64 @@ +#include "../node_context.h" +#include "../op_table.h" +#include "../utils.h" + +#include +#include +#include +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_glu_geglu_quick(const NodeContext & context) { + num_inputs_check(context, 1, 2); + + ov::Output src0; + ov::Output src1; + if (context.get_input_size() == 2) { + src0 = process_view_input_new(context, 0); + src1 = process_view_input_new(context, 1); + } else { + // split along last axis, nc = ne[0] / 2 + auto combined = process_view_input_new(context, 0); + auto combined_shape = combined.get_partial_shape(); + int64_t last_dim_val = combined_shape[combined_shape.rank().get_length() - 1].get_length(); + int64_t nc = last_dim_val / 2; + + auto axis = ov::op::v0::Constant::create(ov::element::i64, {1}, {-1}); + auto step = ov::op::v0::Constant::create(ov::element::i64, {1}, {1}); + auto start0 = ov::op::v0::Constant::create(ov::element::i64, {1}, {0}); + auto stop0 = ov::op::v0::Constant::create(ov::element::i64, {1}, {nc}); + auto start1 = ov::op::v0::Constant::create(ov::element::i64, {1}, {nc}); + auto stop1 = ov::op::v0::Constant::create(ov::element::i64, {1}, {2 * nc}); + + src0 = std::make_shared(combined, start0, stop0, step, axis); + src1 = std::make_shared(combined, start1, stop1, step, axis); + } + + int32_t * params = context.get_output_op_params(); + const int32_t swapped = params[1]; + if (swapped) { + std::swap(src0, src1); + } + + // GELU_QUICK(x) = x * sigmoid(1.702 * x) + // Create the constant in the same type as src0 to avoid f16/f32 mismatch. + auto input_type = src0.get_element_type(); + auto coef = ov::op::v0::Constant::create(input_type, ov::Shape{}, {1.702f}); + auto scaled = std::make_shared(src0, coef); + auto sigmoid = std::make_shared(scaled); + auto gated = std::make_shared(src0, sigmoid); + auto res = std::make_shared(gated, src1); + + return rename_outputs_with_suffix({res}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/pool_2d.cpp b/ggml/src/ggml-openvino/openvino/op/pool_2d.cpp new file mode 100644 index 00000000000..fb6333175f0 --- /dev/null +++ b/ggml/src/ggml-openvino/openvino/op/pool_2d.cpp @@ -0,0 +1,53 @@ +#include "../node_context.h" +#include "../op_table.h" +#include "../utils.h" + +#include +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_pool_2d(const NodeContext & context) { + num_inputs_check(context, 1, 1); + const int32_t * params = context.get_output_op_params(); + + const int k0 = params[1]; + const int k1 = params[2]; + const int s0 = params[3]; + const int s1 = params[4]; + const int p0 = params[5]; + const int p1 = params[6]; + + const int op_case = context.get_op_case(); + ov::Output input = context.get_input(0); + ov::Strides strides{static_cast(s1), static_cast(s0)}; + ov::Shape pads_begin{static_cast(p1), static_cast(p0)}; + ov::Shape pads_end{static_cast(p1), static_cast(p0)}; + ov::Shape kernel{static_cast(k1), static_cast(k0)}; + ov::Output res; + + switch (op_case) { + case 1: // GGML_OP_POOL_MAX + { + res = std::make_shared(input, strides, pads_begin, pads_end, kernel); + break; + } + case 2: // GGML_OP_POOL_AVG + { + res = std::make_shared(input, strides, pads_begin, pads_end, kernel, false); + break; + } + default: + break; + } + return rename_outputs_with_suffix({res}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/roll.cpp b/ggml/src/ggml-openvino/openvino/op/roll.cpp new file mode 100644 index 00000000000..e8d1b8e50b3 --- /dev/null +++ b/ggml/src/ggml-openvino/openvino/op/roll.cpp @@ -0,0 +1,36 @@ +#include "../node_context.h" +#include "../op_table.h" +#include "../utils.h" + +#include +#include + +namespace ov { +namespace frontend { +namespace ggml { +namespace op { + +OutputVector translate_roll(const NodeContext & context) { + num_inputs_check(context, 1, 1); + const int32_t * params = context.get_output_op_params(); + + int64_t s0 = params[0]; + int64_t s1 = params[1]; + int64_t s2 = params[2]; + int64_t s3 = params[3]; + + auto input = context.get_input(0); + + auto shift = ov::op::v0::Constant::create( + ov::element::i64, ov::Shape{4}, std::vector{s3, s2, s1, s0}); + auto axes = ov::op::v0::Constant::create( + ov::element::i64, ov::Shape{4}, std::vector{0, 1, 2, 3}); + + auto roll = std::make_shared(input, shift, axes); + return rename_outputs_with_suffix({roll}, context.get_name()); +} + +} // namespace op +} // namespace ggml +} // namespace frontend +} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/op/view.cpp b/ggml/src/ggml-openvino/openvino/op/view.cpp index 138526cb49c..56f5ceec9bb 100644 --- a/ggml/src/ggml-openvino/openvino/op/view.cpp +++ b/ggml/src/ggml-openvino/openvino/op/view.cpp @@ -17,6 +17,13 @@ namespace op { OutputVector translate_view(const NodeContext & context) { num_inputs_check(context, 1, 1); + if (context.get_op_case() == 1) { + // Static-mode identity pass-through for VIEWs over a GATED_DELTA_NET combined output or + // the conv_input CONCAT; the consuming op (CPY/RMS_NORM) does its own runtime-correct + // slicing on the full tensor (see ggml-decoder.cpp compute_op_case, GGML_OP_VIEW). + return {context.get_input(0)}; + } + if (!context.is_static()) { // On the stateless/non-static path VIEW is normally a no-op (consumers re-slice). // EXCEPTION: the MoE expert aggregation slices each expert plane out of diff --git a/ggml/src/ggml-openvino/openvino/op_table.cpp b/ggml/src/ggml-openvino/openvino/op_table.cpp index 3c26fe83b1a..9c9d8eeac78 100644 --- a/ggml/src/ggml-openvino/openvino/op_table.cpp +++ b/ggml/src/ggml-openvino/openvino/op_table.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -55,10 +56,12 @@ std::unordered_map get_supported_ops() { {"GGML_UNARY_OP_SIGMOID", op::translate_1to1_match_1_input }, {"GGML_UNARY_OP_EXP", op::translate_1to1_match_1_input }, {"GGML_UNARY_OP_NEG", op::translate_1to1_match_1_input }, + {"GGML_UNARY_OP_RELU", op::translate_1to1_match_1_input }, {"GGML_OP_VIEW", op::translate_view }, {"GGML_GLU_OP_SWIGLU", op::translate_glu_swiglu }, {"GGML_GLU_OP_SWIGLU_OAI", op::translate_glu_swiglu_oai }, {"GGML_GLU_OP_GEGLU", op::translate_glu_geglu }, + {"GGML_GLU_OP_GEGLU_QUICK", op::translate_glu_geglu_quick }, {"GGML_OP_SET_ROWS", op::translate_set_rows }, {"GGML_OP_CPY", op::translate_cpy }, {"GGML_OP_FLASH_ATTN_EXT", op::translate_flash_attn_ext }, @@ -72,6 +75,8 @@ std::unordered_map get_supported_ops() { {"GGML_OP_DIAG", op::translate_diag }, {"GGML_OP_TRI", op::translate_tri }, {"GGML_OP_SET", op::translate_set }, + {"GGML_OP_POOL_2D", op::translate_pool_2d }, + {"GGML_OP_ROLL", op::translate_roll }, // solve_tri has accuracy issues on GPU // {"GGML_OP_SOLVE_TRI", op::translate_solve_tri }, }; diff --git a/ggml/src/ggml-openvino/openvino/op_table.h b/ggml/src/ggml-openvino/openvino/op_table.h index d4b9292d637..0a81a57a667 100644 --- a/ggml/src/ggml-openvino/openvino/op_table.h +++ b/ggml/src/ggml-openvino/openvino/op_table.h @@ -38,6 +38,7 @@ GGML_OP_CONVERTER(translate_view); GGML_OP_CONVERTER(translate_glu_swiglu); GGML_OP_CONVERTER(translate_glu_swiglu_oai); GGML_OP_CONVERTER(translate_glu_geglu); +GGML_OP_CONVERTER(translate_glu_geglu_quick); GGML_OP_CONVERTER(translate_set_rows); GGML_OP_CONVERTER(translate_cpy); GGML_OP_CONVERTER(translate_argsort); @@ -53,6 +54,8 @@ GGML_OP_CONVERTER(translate_set); GGML_OP_CONVERTER(translate_diag); GGML_OP_CONVERTER(translate_tri); GGML_OP_CONVERTER(translate_solve_tri); +GGML_OP_CONVERTER(translate_pool_2d); +GGML_OP_CONVERTER(translate_roll); } // namespace op diff --git a/ggml/src/ggml-openvino/openvino/pass/fuse_to_conv.cpp b/ggml/src/ggml-openvino/openvino/pass/fuse_to_conv.cpp new file mode 100644 index 00000000000..21801c0f399 --- /dev/null +++ b/ggml/src/ggml-openvino/openvino/pass/fuse_to_conv.cpp @@ -0,0 +1,212 @@ +#include "fuse_to_conv.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace opp = ov::pass::pattern; + +namespace ov { +namespace frontend { +namespace ggml { +namespace pass { + +// This pass fuses an IM2COL + MatMul convolution into OpenVINO's Convolution op for performance gains. +// Reference the im2col.cpp translator for reference on the pattern being matched. + +FuseToConv::FuseToConv() { + const auto m_wei = opp::any_input(); + const auto m_act = opp::any_input(); + const auto m_matmul = opp::wrap_type({m_wei, m_act}); + + const auto callback = [=](ov::pass::pattern::Matcher & m) { + const auto & pm = m.get_pattern_value_map(); + + auto matmul_node = ov::as_type_ptr(pm.at(m_matmul).get_node_shared_ptr()); + if (!matmul_node || matmul_node->get_transpose_a() || !matmul_node->get_transpose_b()) { + return false; + } + + auto trace = matmul_node->input_value(1); + + // Optional Convert + if (auto n = ov::as_type_ptr(trace.get_node_shared_ptr())) { + trace = n->input_value(0); + } + + for (int i = 0; i < 2; ++i) { + auto n = ov::as_type_ptr(trace.get_node_shared_ptr()); + if (!n) { + return false; + } + trace = n->input_value(0); + } + + if (auto n = ov::as_type_ptr(trace.get_node_shared_ptr())) { + trace = n->input_value(0); + } else { + return false; + } + + if (auto n = ov::as_type_ptr(trace.get_node_shared_ptr())) { + trace = n->input_value(0); + } else { + return false; + } + + if (auto n = ov::as_type_ptr(trace.get_node_shared_ptr())) { + trace = n->input_value(0); + } else { + return false; + } + + auto eip = ov::as_type_ptr(trace.get_node_shared_ptr()); + if (!eip) { + return false; + } + const auto eip_strides = eip->get_strides(); // {stride_h, stride_w} + const auto eip_rates = eip->get_rates(); // {dil_h, dil_w} + + auto pad = ov::as_type_ptr(eip->input_value(0).get_node_shared_ptr()); + if (!pad) { + return false; + } + auto pads_begin_const = + ov::as_type_ptr(pad->input_value(1).get_node_shared_ptr()); + + const auto pads_begin_vals = pads_begin_const->cast_vector(); // {0, 0, pad_h, pad_w} + const std::ptrdiff_t pad_h = static_cast(pads_begin_vals[2]); + const std::ptrdiff_t pad_w = static_cast(pads_begin_vals[3]); + + auto image_input = pad->input_value(0); // [N, IC, 1, IW] NCHW + + auto w_trace = matmul_node->input_value(0); + if (auto n = ov::as_type_ptr(w_trace.get_node_shared_ptr())) { + w_trace = n->input_value(0); + } + for (int i = 0; i < 2; ++i) { + auto n = ov::as_type_ptr(w_trace.get_node_shared_ptr()); + if (!n) { + break; + } + w_trace = n->input_value(0); + } + + auto weight_const = ov::as_type_ptr(w_trace.get_node_shared_ptr()); + if (!weight_const) { + return false; + } + + // Reshape weight to [OC, IC, 1, KW] (OIHW). + const auto w_shape = weight_const->get_shape(); + ov::Shape conv_w_shape; + if (w_shape.size() == 3) { + conv_w_shape = {w_shape[0], w_shape[1], 1, w_shape[2]}; + } else if (w_shape.size() == 4) { + conv_w_shape = {w_shape[1], w_shape[2], 1, w_shape[3]}; + } else { + return false; + } + + auto weight_reshaped = register_new_node(weight_const->get_element_type(), conv_w_shape, + weight_const->get_data_ptr()); + + ov::Output weight_input = weight_reshaped; + if (weight_reshaped->get_element_type() != image_input.get_element_type()) { + weight_input = register_new_node(weight_reshaped, image_input.get_element_type()); + } + + auto conv = register_new_node( + image_input, weight_input, + ov::Strides{static_cast(eip_strides[0]), static_cast(eip_strides[1])}, + ov::CoordinateDiff{pad_h, pad_w}, ov::CoordinateDiff{pad_h, pad_w}, + ov::Strides{static_cast(eip_rates[0]), static_cast(eip_rates[1])}, + ov::op::PadType::EXPLICIT); + + constexpr auto target_type = ov::element::f32; + ov::Output conv_out = conv; + if (conv_out.get_element_type() != target_type) { + conv_out = register_new_node(conv_out, target_type); + } + + std::shared_ptr add_node; + ov::Output bias_input; + for (const auto & consumer_in : matmul_node->output(0).get_target_inputs()) { + auto cast = ov::as_type_ptr(consumer_in.get_node()->shared_from_this()); + if (!cast) { + continue; + } + for (const auto & add_in : cast->output(0).get_target_inputs()) { + auto add = ov::as_type_ptr(add_in.get_node()->shared_from_this()); + if (!add) { + continue; + } + for (size_t i = 0; i < 2; ++i) { + if (ov::as_type_ptr(add->input_value(i).get_node_shared_ptr())) { + bias_input = add->input_value(i); + add_node = add; + break; + } + } + if (add_node) { + break; + } + } + if (add_node) { + break; + } + } + + ov::Output final_out; + std::shared_ptr target_node; + + if (add_node) { + // Reshape bias [OC, 1] → [1, OC, 1, 1] for NCHW broadcasting. + ov::Output bias = bias_input; + if (bias.get_element_type() != target_type) { + bias = register_new_node(bias, target_type); + } + const auto oc = static_cast(conv_w_shape[0]); + auto bias_shape = register_new_node(ov::element::i64, ov::Shape{4}, + std::vector{1, oc, 1, 1}); + bias = register_new_node(bias, bias_shape, false); + final_out = register_new_node(conv_out, bias); + target_node = add_node; + } else { + final_out = conv_out; + target_node = matmul_node; + } + + // Reshape final output back to the target node's original shape if needed. + auto orig_shape = target_node->get_output_partial_shape(0); + if (orig_shape.is_static() && final_out.get_partial_shape() != orig_shape) { + auto shape_const = register_new_node(ov::element::i64, ov::Shape{orig_shape.size()}, + orig_shape.to_shape()); + final_out = register_new_node(final_out, shape_const, false); + } + + final_out.get_node_shared_ptr()->set_friendly_name(target_node->get_friendly_name()); + ov::copy_runtime_info(m.get_matched_nodes(), final_out.get_node_shared_ptr()); + ov::replace_node(target_node, final_out.get_node_shared_ptr()); + + return true; + }; + + register_matcher(std::make_shared(m_matmul, "ov::frontend::ggml::pass::FuseToConv"), callback); +} + +} // namespace pass +} // namespace ggml +} // namespace frontend +} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/pass/fuse_to_conv.h b/ggml/src/ggml-openvino/openvino/pass/fuse_to_conv.h new file mode 100644 index 00000000000..feac14b13ff --- /dev/null +++ b/ggml/src/ggml-openvino/openvino/pass/fuse_to_conv.h @@ -0,0 +1,17 @@ +#include "openvino/pass/matcher_pass.hpp" + +namespace ov { +namespace frontend { +namespace ggml { +namespace pass { + +class FuseToConv : public ov::pass::MatcherPass { +public: + OPENVINO_MATCHER_PASS_RTTI("ov::frontend::ggml::pass::FuseToConv") + FuseToConv(); +}; + +} // namespace pass +} // namespace ggml +} // namespace frontend +} // namespace ov diff --git a/ggml/src/ggml-openvino/openvino/translate_session.cpp b/ggml/src/ggml-openvino/openvino/translate_session.cpp index 35598aba6be..df3a72f3286 100644 --- a/ggml/src/ggml-openvino/openvino/translate_session.cpp +++ b/ggml/src/ggml-openvino/openvino/translate_session.cpp @@ -5,6 +5,7 @@ #include "ggml-openvino/openvino/node_context.h" #include "ggml-openvino/openvino/utils.h" #include "input_model.h" +#include "pass/fuse_to_conv.h" #include "pass/mark_decompression_convert_constant_folding.h" #include "pass/mark_dequantization_subgraph.h" #include "pass/squeeze_matmul.h" @@ -109,7 +110,8 @@ ov::pass::MakeStateful::ParamResPairs get_kv_param_res_pairs( void add_sliced_mask_stateful(TensorMap & tensor_map) { auto create_sliced_mask = [&](const std::string & mask_name, const std::string & sliced_name) { if ((tensor_map.find(mask_name) != tensor_map.end()) && - (tensor_map.find("token_len_per_seq") != tensor_map.end())) { + (tensor_map.find("token_len_per_seq") != tensor_map.end()) && + (tensor_map.find("inp_pos") != tensor_map.end())) { auto token_len_per_seq = tensor_map.at("token_len_per_seq").get_node_shared_ptr(); auto mask = tensor_map.at(mask_name).get_node_shared_ptr(); std::shared_ptr mask_sliced = mask; @@ -137,6 +139,7 @@ void add_sliced_mask_stateful(TensorMap & tensor_map) { }; create_sliced_mask("self_kq_mask", "KQ_mask_sliced"); + create_sliced_mask("KQ_mask", "KQ_mask_sliced"); create_sliced_mask("self_kq_mask_swa", "KQ_mask_swa_sliced"); } @@ -395,6 +398,7 @@ std::shared_ptr TranslateSession::apply_transformations(std::shared_ptr( std::vector{ov::element::u8, ov::element::i8, ov::element::u4, ov::element::i4}); + manager.register_pass(); if (ggml_model_decoder->is_stateful()) { const auto kv_param_res_names = ggml_model_decoder->get_kv_param_res_names(); diff --git a/ggml/src/ggml-openvino/openvino/utils.cpp b/ggml/src/ggml-openvino/openvino/utils.cpp index 504d74b7067..8bb7678ee38 100644 --- a/ggml/src/ggml-openvino/openvino/utils.cpp +++ b/ggml/src/ggml-openvino/openvino/utils.cpp @@ -72,6 +72,7 @@ OutputVector rename_outputs_with_suffix(const OutputVector & outputs, const std: name += "_"; name += suffix; node->set_friendly_name(name); + // Uncomment to dump every node's inferred shape (used to hunt down dynamic dims on NPU). // std::cout << name << " " << output.get_partial_shape() << std::endl; } return outputs; diff --git a/ggml/src/ggml-openvino/utils.cpp b/ggml/src/ggml-openvino/utils.cpp index 4df8381dcbd..93b1ccbe907 100644 --- a/ggml/src/ggml-openvino/utils.cpp +++ b/ggml/src/ggml-openvino/utils.cpp @@ -16,6 +16,8 @@ #include #include #include +#include +#include #include #include #include @@ -48,7 +50,7 @@ enum ggml_status ov_graph_compute(ggml_cgraph * cgraph, ggml_backend_t backend) GgmlOvDecoder::dump_cgraph(cgraph, filename); } - const auto is_static = ggml_openvino_is_npu(); + const auto is_static = ggml_openvino_is_npu() || ggml_openvino_getenv_int("GGML_OPENVINO_FORCE_STATIC"); GGML_ASSERT(ctx->runtime_context != nullptr); std::shared_ptr r_ctx = std::static_pointer_cast(ctx->runtime_context); @@ -168,13 +170,24 @@ ov::Tensor create_ov_output_tensor(std::shared_ptr ggml_decoder, auto output_type = ggml_decoder->get_ov_type(ggml_tensor); ov::Shape output_shape; + void * output_data = ggml_tensor->data; if (ggml_decoder->is_static()) { output_shape = infer_request->get_output_tensor(output_index).get_shape(); } else { - output_shape = ggml_decoder->get_shape(ggml_tensor); + // For a CPY into a padded view_src (e.g. a padded KV cache buffer), the + // OV ScatterUpdate node outputs the full view_src shape, not the CPY node's + // own (smaller) shape. Using the CPY shape here causes set_output_tensor to + // fail with a shape-incompatibility error. Use view_src's shape and data + // pointer instead so the OV tensor matches the model output exactly. + if (ggml_tensor->op == GGML_OP_CPY && ggml_tensor->view_src != nullptr && + ggml_nbytes(ggml_tensor) != ggml_nbytes(ggml_tensor->view_src)) { + output_shape = ggml_decoder->get_shape(ggml_tensor->view_src); + output_data = ggml_tensor->view_src->data; + } else { + output_shape = ggml_decoder->get_shape(ggml_tensor); + } } - - ov::Tensor output_tensor(output_type, output_shape, ggml_tensor->data); + ov::Tensor output_tensor(output_type, output_shape, output_data); return output_tensor; } @@ -583,7 +596,9 @@ enum ggml_status ov_graph_compute_static(ggml_cgraph * cgraph, std::shared_ptr(ggml_decoder_prefill); - auto input_model_decode = std::make_shared(ggml_decoder_decode); - - auto model_prefill = ov::frontend::ggml::FrontEnd::convert(input_model_prefill); - ggml_decoder_prefill->clear_model_weights(); - auto model_decode = ov::frontend::ggml::FrontEnd::convert(input_model_decode); - ggml_decoder_decode->clear_model_weights(); - conversion_end_time = ggml_time_us(); - - if (ggml_openvino_getenv_int("GGML_OPENVINO_DUMP_IR")) { - char timestamped_filename[64]; - auto timestamp = (long long) ggml_time_us(); - snprintf(timestamped_filename, sizeof(timestamped_filename), "model_prefill_%lld.xml", timestamp); - ov::serialize(model_prefill, timestamped_filename); - snprintf(timestamped_filename, sizeof(timestamped_filename), "model_decode_%lld.xml", timestamp); - ov::serialize(model_decode, timestamped_filename); - } + const bool dump_ir = ggml_openvino_getenv_int("GGML_OPENVINO_DUMP_IR"); + const auto dump_ir_timestamp = static_cast(ggml_time_us()); + + auto build_static_model = [&core, &config, dump_ir, dump_ir_timestamp]( + std::shared_ptr decoder, + const char * tag, + std::shared_ptr & model, + ov::CompiledModel & compiled_model, + std::shared_ptr & infer_request, + int64_t & local_conversion_end_time, + int64_t & local_compile_end_time) { + auto input_model = std::make_shared(decoder); + model = ov::frontend::ggml::FrontEnd::convert(input_model); + decoder->clear_model_weights(); + local_conversion_end_time = ggml_time_us(); + + if (dump_ir) { + char timestamped_filename[64]; + snprintf(timestamped_filename, sizeof(timestamped_filename), "model_%s_%lld.xml", tag, + dump_ir_timestamp); + ov::serialize(model, timestamped_filename); + } + compiled_model = core.compile_model(model, device, config); + infer_request = std::make_shared(compiled_model.create_infer_request()); + local_compile_end_time = ggml_time_us(); + }; + std::shared_ptr model_prefill; + std::shared_ptr model_decode; ov::CompiledModel compiled_model_prefill; ov::CompiledModel compiled_model_decode; - auto remote_context = ggml_openvino_get_remote_context(); - if (remote_context.has_value()) { - compiled_model_prefill = core.compile_model(model_prefill, remote_context.value(), config); - compiled_model_decode = core.compile_model(model_decode, remote_context.value(), config); - } else { - compiled_model_prefill = core.compile_model(model_prefill, device, config); - compiled_model_decode = core.compile_model(model_decode, device, config); - } - - auto infer_request_prefill = std::make_shared(compiled_model_prefill.create_infer_request()); - auto infer_request_decode = std::make_shared(compiled_model_decode.create_infer_request()); - compile_end_time = ggml_time_us(); + std::shared_ptr infer_request_prefill; + std::shared_ptr infer_request_decode; + int64_t prefill_conversion_end_time; + int64_t decode_conversion_end_time; + int64_t prefill_compile_end_time; + int64_t decode_compile_end_time; + auto prefill_future = std::async(std::launch::async, build_static_model, ggml_decoder_prefill, "prefill", + std::ref(model_prefill), std::ref(compiled_model_prefill), + std::ref(infer_request_prefill), std::ref(prefill_conversion_end_time), + std::ref(prefill_compile_end_time)); + auto decode_future = std::async(std::launch::async, build_static_model, ggml_decoder_decode, "decode", + std::ref(model_decode), std::ref(compiled_model_decode), + std::ref(infer_request_decode), std::ref(decode_conversion_end_time), + std::ref(decode_compile_end_time)); + prefill_future.get(); + decode_future.get(); + conversion_end_time = std::max(prefill_conversion_end_time, decode_conversion_end_time); + compile_end_time = std::max(prefill_compile_end_time, decode_compile_end_time); model = is_prefill ? model_prefill : model_decode; ggml_decoder = is_prefill ? ggml_decoder_prefill : ggml_decoder_decode; @@ -742,7 +774,7 @@ enum ggml_status ov_graph_compute_static(ggml_cgraph * cgraph, std::shared_ptrne[0]; + auto inp_len = get_inp_pos_n_tokens(cgraph, inp_pos); for (int chunk_index = 0; chunk_index * prefill_chunk_size < inp_len; chunk_index++) { for (size_t i = 0; i < ov_input_names_local.size(); i++) { auto param_name = ov_input_names_local[i]; @@ -762,6 +794,11 @@ enum ggml_status ov_graph_compute_static(ggml_cgraph * cgraph, std::shared_ptrsecond; + if (ggml_nbytes(ggml_tensor) == 0) { + // Zero-row in-place writeback (e.g. the empty s_copy defrag remainder). The OV + // Result is the full cache, so binding it over this 0-byte buffer overflows it. + continue; + } auto output_tensor = create_ov_output_tensor(ggml_decoder, infer_request, i, ggml_tensor); infer_request->set_output_tensor(i, output_tensor); } @@ -798,6 +835,9 @@ enum ggml_status ov_graph_compute_static(ggml_cgraph * cgraph, std::shared_ptrsecond; + if (ggml_nbytes(ggml_tensor) == 0) { + continue; + } auto output_tensor = create_ov_output_tensor(ggml_decoder, infer_request, i, ggml_tensor); infer_request->set_output_tensor(i, output_tensor); } @@ -1074,6 +1114,9 @@ ov::Tensor get_ov_input_tensor(std::shared_ptr ggml_decoder, cons ov::Tensor get_ov_input_tensor_static_decode(std::shared_ptr ggml_decoder, const std::string & param_name) { // NPU decoding stage + if (ggml_decoder->get_model_extra_inputs().count(param_name)) { + return get_ov_input_tensor(ggml_decoder, param_name); + } const auto * ggml_tensor = ggml_decoder->get_input_ggml_tensor(param_name); const auto * op = ggml_decoder->get_tensor_used_op(ggml_tensor); @@ -1123,14 +1166,30 @@ ov::Tensor get_ov_input_tensor_static_prefill(std::shared_ptr ggm const std::string & param_name, int chunk_index) { // NPU prompt processing stage - const auto * ggml_tensor = ggml_decoder->get_input_ggml_tensor(param_name); - const auto * op = ggml_decoder->get_tensor_used_op(ggml_tensor); - const size_t input_len = ggml_decoder->get_input_len(); const size_t chunk_size = ggml_decoder->m_prefill_chunk_size; const size_t chunk_valid_size = std::min(chunk_size, input_len - chunk_index * chunk_size); const size_t chunk_pad_size = chunk_size - chunk_valid_size; + if (param_name == "chunk_valid_len") { + ov::Tensor input_tensor(ov::element::i64, ov::Shape{1}); + *input_tensor.data() = (int64_t) chunk_valid_size; + return input_tensor; + } + if (chunk_index > 0 && param_name == "cache_rs_reset_len") { + // The recurrent-state clear belongs to the start of the sequence. Re-applying it on every + // chunk would wipe the state accumulated by the preceding chunks, so disable it (a zero + // length makes scale.cpp's keep-mask select every slot) after the first chunk. + ov::Tensor input_tensor(ov::element::i64, ov::Shape{1}); + *input_tensor.data() = 0; + return input_tensor; + } + if (ggml_decoder->get_model_extra_inputs().count(param_name)) { + return get_ov_input_tensor(ggml_decoder, param_name); + } + const auto * ggml_tensor = ggml_decoder->get_input_ggml_tensor(param_name); + const auto * op = ggml_decoder->get_tensor_used_op(ggml_tensor); + if (GgmlOvDecoder::is_inp_pos(ggml_tensor, op) && GgmlOvDecoder::get_inp_pos_n_planes(op) > 1) { // IMROPE: inp_pos stacks n_planes (t/h/w/e) position planes, each of length // input_len; pad every plane independently so they stay aligned to chunk_size. @@ -1306,7 +1365,7 @@ void print_input_tensor_info(const std::string & name, const ov::Tensor & tensor << std::endl; switch (tensor.get_element_type()) { case ov::element::f32: { - if (name.find("self_kq_mask") == std::string::npos) { + if (name.find("self_kq_mask") == std::string::npos && name.find("KQ_mask") == std::string::npos) { std::cout << *(tensor.data()) << std::endl; } else { size_t rows = tensor.get_shape()[2]; @@ -1414,8 +1473,24 @@ const ggml_tensor * get_inp_pos_tensor(ggml_cgraph * cgraph) { throw std::runtime_error("get_inp_pos_tensor: inp_pos not found in cgraph"); } -bool get_is_prefill(const ggml_tensor * inp_pos) { - return inp_pos->ne[0] > 1; +int64_t get_inp_pos_n_tokens(ggml_cgraph * cgraph, const ggml_tensor * inp_pos) { + // IMROPE stacks n_planes (t/h/w/e) position planes into inp_pos, so ne[0] is + // n_planes * n_tokens. Callers that need a token count must divide the planes out. + int n_planes = 1; + for (int i = 0; i < cgraph->n_nodes; ++i) { + auto * op = cgraph->nodes[i]; + for (int j = 0; j < GGML_MAX_SRC; ++j) { + if (op->src[j] == inp_pos) { + n_planes = GgmlOvDecoder::get_inp_pos_n_planes(op); + break; + } + } + } + return inp_pos->ne[0] / n_planes; +} + +bool get_is_prefill(ggml_cgraph * cgraph, const ggml_tensor * inp_pos) { + return get_inp_pos_n_tokens(cgraph, inp_pos) > 1; } #pragma GCC diagnostic pop diff --git a/ggml/src/ggml-openvino/utils.h b/ggml/src/ggml-openvino/utils.h index 513fa83c9d6..5aa74da38d3 100644 --- a/ggml/src/ggml-openvino/utils.h +++ b/ggml/src/ggml-openvino/utils.h @@ -164,7 +164,9 @@ std::vector pad_input(const ggml_tensor * tensor, size_t padded_rows, size_t const ggml_tensor * get_inp_pos_tensor(struct ggml_cgraph * cgraph); -bool get_is_prefill(const ggml_tensor * inp_pos); +int64_t get_inp_pos_n_tokens(struct ggml_cgraph * cgraph, const ggml_tensor * inp_pos); + +bool get_is_prefill(struct ggml_cgraph * cgraph, const ggml_tensor * inp_pos); ov::Tensor get_ov_input_tensor(std::shared_ptr ggml_decoder, const std::string & param_name); ov::Tensor get_ov_input_tensor_static_decode(std::shared_ptr ggml_decoder, diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 69a8a08ae17..9aa5883d80d 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -9,6 +9,9 @@ #include #include #include +#include +#include +#include #include #include #include @@ -17,6 +20,8 @@ #include #include #include +#include +#include static const char * RPC_DEBUG = std::getenv("GGML_RPC_DEBUG"); @@ -72,6 +77,7 @@ enum rpc_cmd { RPC_CMD_DEVICE_COUNT, RPC_CMD_GRAPH_RECOMPUTE, RPC_CMD_MEMSET_TENSOR, + RPC_CMD_NONE, RPC_CMD_COUNT, }; @@ -223,24 +229,24 @@ struct ggml_backend_rpc_buffer_type_context { size_t max_size; }; +class rpc_dispatcher; struct ggml_backend_rpc_context { - std::string endpoint; - uint32_t device; - std::string name; + std::shared_ptr dispatcher; + uint32_t device; + std::string name; }; struct ggml_backend_rpc_buffer_context { - std::shared_ptr sock; - void * base_ptr; - uint64_t remote_ptr; + std::shared_ptr dispatcher; + void * base_ptr; + uint64_t remote_ptr; }; // RPC helper functions // Computes FNV-1a hash of the data -static uint64_t fnv_hash(const uint8_t * data, size_t len) { +static uint64_t fnv_hash(const uint8_t * data, size_t len, uint64_t hash = 0xcbf29ce484222325ULL) { const uint64_t fnv_prime = 0x100000001b3ULL; - uint64_t hash = 0xcbf29ce484222325ULL; for (size_t i = 0; i < len; ++i) { hash ^= data[i]; @@ -357,44 +363,248 @@ static bool negotiate_hello(const std::shared_ptr & sock) { return true; } -static std::shared_ptr get_socket(const std::string & endpoint) { - static std::mutex mutex; - std::lock_guard lock(mutex); - static std::unordered_map> sockets; +template +class message_queue { +public: + message_queue() {} - auto it = sockets.find(endpoint); - if (it != sockets.end()) { - if (auto sock = it->second.lock()) { - return sock; + bool push(const T &value) { + std::unique_lock lock(mutex); + if (interrupted) { + return false; } + queue.push(value); + cvar.notify_all(); + return true; } + + bool pop(T* out) { + std::unique_lock lock(mutex); + cvar.wait(lock, [this] { return !queue.empty() || interrupted; }); + if (interrupted) { + return false; + } + *out = queue.front(); + queue.pop(); + return true; + } + + void interrupt() { + std::unique_lock lock(mutex); + interrupted = true; + lock.unlock(); + cvar.notify_all(); + } + +private: + bool interrupted = false; + std::queue queue; + std::mutex mutex; + std::condition_variable cvar; +}; + +class rpc_dispatcher { +public: + rpc_dispatcher() { + } + + void send(enum rpc_cmd cmd, std::shared_ptr input, size_t input_size); + void send(enum rpc_cmd cmd, std::shared_ptr input, size_t input_size, void * output, size_t output_size); + void send_async(enum rpc_cmd cmd, std::shared_ptr input, size_t input_size); + void send_async(enum rpc_cmd cmd, std::shared_ptr input, size_t input_size, void * output, size_t output_size); + + ggml_backend_event_t event_new(ggml_backend_dev_t dev); + void event_free(ggml_backend_event_t event); + void event_synchronize(ggml_backend_event_t event); + void event_record(ggml_backend_event_t event); + void synchronize(); + + void start(const std::string & endpoint); + void work(); + + ~rpc_dispatcher(); + +private: + struct rpc_msg { + rpc_cmd cmd; + std::shared_ptr input; + size_t input_size; + void * output; + size_t output_size; + std::promise completion; + }; + using rpc_msg_ptr = std::shared_ptr; + using rpc_msg_queue = message_queue; + struct rpc_event { + rpc_msg_ptr msg; + std::shared_future sf; + }; + rpc_msg_queue queue; + socket_ptr sock; + std::atomic_bool running; + std::thread thread; +}; + +static void rpc_dispatcher_trampoline(rpc_dispatcher * dispatcher) +{ + dispatcher->work(); +} + +void rpc_dispatcher::send(enum rpc_cmd cmd, std::shared_ptr input, size_t input_size) { + auto msg = std::make_shared(); + msg->cmd = cmd; + msg->input = input; + msg->input_size = input_size; + msg->output = nullptr; + msg->output_size = 0; + GGML_ASSERT(queue.push(msg)); + auto future = msg->completion.get_future(); + future.wait(); +} + +void rpc_dispatcher::send_async(enum rpc_cmd cmd, std::shared_ptr input, size_t input_size) { + auto msg = std::make_shared(); + msg->cmd = cmd; + msg->input = input; + msg->input_size = input_size; + msg->output = nullptr; + msg->output_size = 0; + GGML_ASSERT(queue.push(msg)); +} + +void rpc_dispatcher::send(enum rpc_cmd cmd, std::shared_ptr input, size_t input_size, void * output, size_t output_size) { + auto msg = std::make_shared(); + msg->cmd = cmd; + msg->input = input; + msg->input_size = input_size; + msg->output = output; + msg->output_size = output_size; + GGML_ASSERT(queue.push(msg)); + auto future = msg->completion.get_future(); + future.wait(); +} + +void rpc_dispatcher::send_async(enum rpc_cmd cmd, std::shared_ptr input, size_t input_size, void * output, size_t output_size) { + auto msg = std::make_shared(); + msg->cmd = cmd; + msg->input = input; + msg->input_size = input_size; + msg->output = output; + msg->output_size = output_size; + GGML_ASSERT(queue.push(msg)); +} + +ggml_backend_event_t rpc_dispatcher::event_new(ggml_backend_dev_t dev) { + rpc_event * ev = new rpc_event; + ev->msg = std::make_shared(); + ev->msg->cmd = RPC_CMD_NONE; + ev->sf = ev->msg->completion.get_future().share(); + GGML_ASSERT(queue.push(ev->msg)); + return new ggml_backend_event { + /* .device = */ dev, + /* .context = */ ev, + }; +} + +void rpc_dispatcher::event_free(ggml_backend_event_t event) { + rpc_event * ev = (rpc_event *)event->context; + delete ev; +} + +void rpc_dispatcher::event_synchronize(ggml_backend_event_t event) { + rpc_event * ev = (rpc_event *)event->context; + ev->sf.wait(); +} + +void rpc_dispatcher::event_record(ggml_backend_event_t event) { + rpc_event * ev = (rpc_event *)event->context; + ev->msg = std::make_shared(); + ev->msg->cmd = RPC_CMD_NONE; + ev->sf = ev->msg->completion.get_future().share(); + GGML_ASSERT(queue.push(ev->msg)); +} + +void rpc_dispatcher::synchronize() { + // to ensure all messages are processed, submit dummy message and wait for it to complete + auto msg = std::make_shared(); + msg->cmd = RPC_CMD_NONE; + GGML_ASSERT(queue.push(msg)); + msg->completion.get_future().wait(); +} + +void rpc_dispatcher::start(const std::string & endpoint) { std::string host; int port; if (!parse_endpoint(endpoint, host, port)) { - GGML_LOG_ERROR("Failed to parse endpoint: %s\n", endpoint.c_str()); - return nullptr; + GGML_ABORT("Failed to parse endpoint: %s\n", endpoint.c_str()); } - if (!rpc_transport_init()) { - return nullptr; + GGML_ABORT("RPC transport initialization failed\n"); } - auto sock = socket_t::connect(host.c_str(), port); + + sock = socket_t::connect(host.c_str(), port); if (sock == nullptr) { - return nullptr; + GGML_ABORT("Failed to connect to %s\n", endpoint.c_str()); } if (!negotiate_hello(sock)) { - return nullptr; + GGML_ABORT("RPC handshake failed for %s\n", endpoint.c_str()); } LOG_DBG("[%s] connected to %s\n", __func__, endpoint.c_str()); - sockets[endpoint] = sock; - return sock; + running = true; + thread = std::thread(rpc_dispatcher_trampoline, this); +} + +void rpc_dispatcher::work() { + while (running) { + rpc_msg_ptr msg_ptr; + if (!queue.pop(&msg_ptr)) { + break; + } + if (msg_ptr->cmd != RPC_CMD_NONE) { + if (msg_ptr->output) { + bool status = send_rpc_cmd(sock, msg_ptr->cmd, msg_ptr->input.get(), msg_ptr->input_size, msg_ptr->output, msg_ptr->output_size); + RPC_STATUS_ASSERT(status); + } else { + bool status = send_rpc_cmd(sock, msg_ptr->cmd, msg_ptr->input.get(), msg_ptr->input_size); + RPC_STATUS_ASSERT(status); + } + } + msg_ptr->completion.set_value(); + } +} + +rpc_dispatcher::~rpc_dispatcher() { + running = false; + queue.interrupt(); + sock = nullptr; + if (thread.joinable()) { + thread.join(); + } +} + +static std::shared_ptr get_dispatcher(const std::string & endpoint) { + static std::mutex mutex; + std::lock_guard lock(mutex); + static std::unordered_map> dispatchers; + + auto it = dispatchers.find(endpoint); + if (it != dispatchers.end()) { + if (auto dispatcher = it->second.lock()) { + return dispatcher; + } + } + + auto dispatcher = std::make_shared(); + dispatcher->start(endpoint); + dispatchers[endpoint] = dispatcher; + return dispatcher; } static void ggml_backend_rpc_buffer_free_buffer(ggml_backend_buffer_t buffer) { ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context; - rpc_msg_free_buffer_req request = {ctx->remote_ptr}; - bool status = send_rpc_cmd(ctx->sock, RPC_CMD_FREE_BUFFER, &request, sizeof(request), nullptr, 0); - RPC_STATUS_ASSERT(status); + auto request = std::make_shared(); + request->remote_ptr = ctx->remote_ptr; + ctx->dispatcher->send(RPC_CMD_FREE_BUFFER, request, sizeof(*request)); delete ctx; } @@ -403,10 +613,10 @@ static void * ggml_backend_rpc_buffer_get_base(ggml_backend_buffer_t buffer) { if (ctx->base_ptr != nullptr) { return ctx->base_ptr; } - rpc_msg_buffer_get_base_req request = {ctx->remote_ptr}; + auto request = std::make_shared(); + request->remote_ptr = ctx->remote_ptr; rpc_msg_buffer_get_base_rsp response; - bool status = send_rpc_cmd(ctx->sock, RPC_CMD_BUFFER_GET_BASE, &request, sizeof(request), &response, sizeof(response)); - RPC_STATUS_ASSERT(status); + ctx->dispatcher->send(RPC_CMD_BUFFER_GET_BASE, request, sizeof(*request), &response, sizeof(response)); ctx->base_ptr = reinterpret_cast(response.base_ptr); return ctx->base_ptr; } @@ -463,12 +673,9 @@ static enum ggml_status ggml_backend_rpc_buffer_init_tensor(ggml_backend_buffer_ // Due to bandwidth constraints, we only call the server init tensor functions if necessary. // In particular, only quantized tensors need padding if (ggml_is_quantized(tensor->type) && (tensor->ne[0] % 512 != 0) && (tensor->view_src == nullptr)) { - rpc_msg_init_tensor_req request; - - request.tensor = serialize_tensor(tensor); - - bool status = send_rpc_cmd(ctx->sock, RPC_CMD_INIT_TENSOR, &request, sizeof(request), nullptr, 0); - RPC_STATUS_ASSERT(status); + auto request = std::make_shared(); + request->tensor = serialize_tensor(tensor); + ctx->dispatcher->send(RPC_CMD_INIT_TENSOR, request, sizeof(*request)); } return GGML_STATUS_SUCCESS; } @@ -476,27 +683,24 @@ static enum ggml_status ggml_backend_rpc_buffer_init_tensor(ggml_backend_buffer_ static void ggml_backend_rpc_buffer_memset_tensor( ggml_backend_buffer_t buffer, ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) { ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context; - rpc_msg_memset_tensor_req request = { - /* .tensor = */ serialize_tensor(tensor), - /* .offset = */ offset, - /* .size = */ size, - /* .value = */ value, - }; - bool status = send_rpc_cmd(ctx->sock, RPC_CMD_MEMSET_TENSOR, &request, sizeof(request), nullptr, 0); - RPC_STATUS_ASSERT(status); + auto request = std::make_shared(); + request->tensor = serialize_tensor(tensor); + request->offset = offset; + request->size = size; + request->value = value; + ctx->dispatcher->send(RPC_CMD_MEMSET_TENSOR, request, sizeof(*request)); } static void ggml_backend_rpc_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context; rpc_tensor rpc_tensor = serialize_tensor(tensor); if (size > HASH_THRESHOLD) { - rpc_msg_set_tensor_hash_req request; - request.tensor = rpc_tensor; - request.offset = offset; - request.hash = fnv_hash((const uint8_t*)data, size); + auto request = std::make_shared(); + request->tensor = rpc_tensor; + request->offset = offset; + request->hash = fnv_hash((const uint8_t*)data, size); rpc_msg_set_tensor_hash_rsp response; - bool status = send_rpc_cmd(ctx->sock, RPC_CMD_SET_TENSOR_HASH, &request, sizeof(request), &response, sizeof(response)); - RPC_STATUS_ASSERT(status); + ctx->dispatcher->send(RPC_CMD_SET_TENSOR_HASH, request, sizeof(*request), &response, sizeof(response)); if (response.result) { // the server has the same data, no need to send it return; @@ -504,22 +708,21 @@ static void ggml_backend_rpc_buffer_set_tensor(ggml_backend_buffer_t buffer, ggm } // input serialization format: | rpc_tensor | offset (8 bytes) | data (size bytes) size_t input_size = sizeof(rpc_tensor) + sizeof(uint64_t) + size; - std::vector input(input_size, 0); - memcpy(input.data(), &rpc_tensor, sizeof(rpc_tensor)); - memcpy(input.data() + sizeof(rpc_tensor), &offset, sizeof(offset)); - memcpy(input.data() + sizeof(rpc_tensor) + sizeof(offset), data, size); - bool status = send_rpc_cmd(ctx->sock, RPC_CMD_SET_TENSOR, input.data(), input.size()); - RPC_STATUS_ASSERT(status); + uint8_t * input = new uint8_t[input_size](); + memcpy(input, &rpc_tensor, sizeof(rpc_tensor)); + memcpy(input + sizeof(rpc_tensor), &offset, sizeof(offset)); + memcpy(input + sizeof(rpc_tensor) + sizeof(offset), data, size); + std::shared_ptr input_ptr(input, std::default_delete()); + ctx->dispatcher->send(RPC_CMD_SET_TENSOR, input_ptr, input_size); } static void ggml_backend_rpc_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context; - rpc_msg_get_tensor_req request; - request.tensor = serialize_tensor(tensor); - request.offset = offset; - request.size = size; - bool status = send_rpc_cmd(ctx->sock, RPC_CMD_GET_TENSOR, &request, sizeof(request), data, size); - RPC_STATUS_ASSERT(status); + auto request = std::make_shared(); + request->tensor = serialize_tensor(tensor); + request->offset = offset; + request->size = size; + ctx->dispatcher->send(RPC_CMD_GET_TENSOR, request, sizeof(*request), data, size); } static bool ggml_backend_rpc_buffer_cpy_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * src, ggml_tensor * dst) { @@ -529,16 +732,15 @@ static bool ggml_backend_rpc_buffer_cpy_tensor(ggml_backend_buffer_t buffer, con ggml_backend_rpc_buffer_context * src_ctx = (ggml_backend_rpc_buffer_context *)src_buffer->context; ggml_backend_buffer_t dst_buffer = dst->buffer; ggml_backend_rpc_buffer_context * dst_ctx = (ggml_backend_rpc_buffer_context *)dst_buffer->context; - if (src_ctx->sock != dst_ctx->sock) { + if (src_ctx->dispatcher != dst_ctx->dispatcher) { return false; } ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context; - rpc_msg_copy_tensor_req request; - request.src = serialize_tensor(src); - request.dst = serialize_tensor(dst); + auto request = std::make_shared(); + request->src = serialize_tensor(src); + request->dst = serialize_tensor(dst); rpc_msg_copy_tensor_rsp response; - bool status = send_rpc_cmd(ctx->sock, RPC_CMD_COPY_TENSOR, &request, sizeof(request), &response, sizeof(response)); - RPC_STATUS_ASSERT(status); + ctx->dispatcher->send(RPC_CMD_COPY_TENSOR, request, sizeof(*request), &response, sizeof(response)); return response.result; } return false; @@ -546,9 +748,10 @@ static bool ggml_backend_rpc_buffer_cpy_tensor(ggml_backend_buffer_t buffer, con static void ggml_backend_rpc_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { ggml_backend_rpc_buffer_context * ctx = (ggml_backend_rpc_buffer_context *)buffer->context; - rpc_msg_buffer_clear_req request = {ctx->remote_ptr, value}; - bool status = send_rpc_cmd(ctx->sock, RPC_CMD_BUFFER_CLEAR, &request, sizeof(request), nullptr, 0); - RPC_STATUS_ASSERT(status); + auto request = std::make_shared(); + request->remote_ptr = ctx->remote_ptr; + request->value = value; + ctx->dispatcher->send(RPC_CMD_BUFFER_CLEAR, request, sizeof(*request)); } static ggml_backend_buffer_i ggml_backend_rpc_buffer_interface = { @@ -572,15 +775,17 @@ static const char * ggml_backend_rpc_buffer_type_name(ggml_backend_buffer_type_t static ggml_backend_buffer_t ggml_backend_rpc_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { ggml_backend_rpc_buffer_type_context * buft_ctx = (ggml_backend_rpc_buffer_type_context *)buft->context; - rpc_msg_alloc_buffer_req request = {buft_ctx->device, size}; + auto request = std::make_shared(); + request->device = buft_ctx->device; + request->size = size; rpc_msg_alloc_buffer_rsp response; - auto sock = get_socket(buft_ctx->endpoint); - bool status = send_rpc_cmd(sock, RPC_CMD_ALLOC_BUFFER, &request, sizeof(request), &response, sizeof(response)); - RPC_STATUS_ASSERT(status); + + auto dispatcher = get_dispatcher(buft_ctx->endpoint); + dispatcher->send(RPC_CMD_ALLOC_BUFFER, request, sizeof(*request), &response, sizeof(response)); if (response.remote_ptr != 0) { ggml_backend_buffer_t buffer = ggml_backend_buffer_init(buft, ggml_backend_rpc_buffer_interface, - new ggml_backend_rpc_buffer_context{sock, nullptr, response.remote_ptr}, + new ggml_backend_rpc_buffer_context{dispatcher, nullptr, response.remote_ptr}, response.remote_size); return buffer; } else { @@ -588,11 +793,11 @@ static ggml_backend_buffer_t ggml_backend_rpc_buffer_type_alloc_buffer(ggml_back } } -static size_t get_alignment(const std::shared_ptr & sock, uint32_t device) { - rpc_msg_get_alignment_req request = {device}; +static size_t get_alignment(const std::shared_ptr & dispatcher, uint32_t device) { + auto request = std::make_shared(); + request->device = device; rpc_msg_get_alignment_rsp response; - bool status = send_rpc_cmd(sock, RPC_CMD_GET_ALIGNMENT, &request, sizeof(request), &response, sizeof(response)); - RPC_STATUS_ASSERT(status); + dispatcher->send(RPC_CMD_GET_ALIGNMENT, request, sizeof(*request), &response, sizeof(response)); return response.alignment; } @@ -601,11 +806,11 @@ static size_t ggml_backend_rpc_buffer_type_get_alignment(ggml_backend_buffer_typ return buft_ctx->alignment; } -static size_t get_max_size(const std::shared_ptr & sock, uint32_t device) { - rpc_msg_get_max_size_req request = {device}; +static size_t get_max_size(const std::shared_ptr & dispatcher, uint32_t device) { + auto request = std::make_shared(); + request->device = device; rpc_msg_get_max_size_rsp response; - bool status = send_rpc_cmd(sock, RPC_CMD_GET_MAX_SIZE, &request, sizeof(request), &response, sizeof(response)); - RPC_STATUS_ASSERT(status); + dispatcher->send(RPC_CMD_GET_MAX_SIZE, request, sizeof(*request), &response, sizeof(response)); return response.max_size; } @@ -628,23 +833,63 @@ static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_ty if (rpc_get) { ggml_backend_rpc_buffer_type_context * buft_ctx = (ggml_backend_rpc_buffer_type_context *)buft->context; - auto sock = get_socket(buft_ctx->endpoint); - rpc_msg_get_alloc_size_req request = { - /*.device =*/ buft_ctx->device, - /*.tensor =*/ serialize_tensor(tensor), - /*.srcs =*/ {}, + // Cache key for calls to read the alloc_size. + // We deliberately exclude src tensor dimensions from the key because: + // 1. For CPU backends, alloc_size = ggml_nbytes(output) regardless of src shapes + // 2. For GPU backends, the reservation graph uses max dimensions, so the + // cached value from reservation is always >= any subsequent request + // 3. Including src dims causes cache misses per-ubatch (e.g. growing KV cache) + // which blocks the main thread behind in-flight GRAPH_COMPUTE commands + struct alloc_size_cache_key { + uint32_t device; + uint32_t type; + uint32_t op; + int32_t op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t)]; + uint32_t ne[GGML_MAX_DIMS]; }; + alloc_size_cache_key key = {}; + key.device = buft_ctx->device; + key.type = tensor->type; + key.op = tensor->op; + memcpy(key.op_params, tensor->op_params, sizeof(key.op_params)); + for (int i = 0; i < GGML_MAX_DIMS; i++) { + key.ne[i] = (uint32_t)tensor->ne[i]; + } + + uint64_t cache_hash = fnv_hash((const uint8_t *)&key, sizeof(key)); + cache_hash = fnv_hash((const uint8_t *)buft_ctx->endpoint.data(), buft_ctx->endpoint.size(), cache_hash); + + // alloc sizes are immutable for a given tensor configuration + static std::mutex cache_mutex; + static std::unordered_map cache; + + { + std::lock_guard lock(cache_mutex); + auto it = cache.find(cache_hash); + if (it != cache.end()) { + return it->second; + } + } + + auto request = std::make_shared(); + request->device = buft_ctx->device; + request->tensor = serialize_tensor(tensor); + // .get_alloc_size could be a function of the tensor's srcs, so we must serialize them as well for (int i = 0; i < GGML_MAX_SRC; i++) { - request.srcs[i] = serialize_tensor(tensor->src[i]); + request->srcs[i] = serialize_tensor(tensor->src[i]); } - // TODO: cache the alloc responses to avoid extra RPC calls? rpc_msg_get_alloc_size_rsp response; - bool status = send_rpc_cmd(sock, RPC_CMD_GET_ALLOC_SIZE, &request, sizeof(request), &response, sizeof(response)); - RPC_STATUS_ASSERT(status); + auto dispatcher = get_dispatcher(buft_ctx->endpoint); + dispatcher->send(RPC_CMD_GET_ALLOC_SIZE, request, sizeof(*request), &response, sizeof(response)); + + { + std::lock_guard lock(cache_mutex); + cache[cache_hash] = response.alloc_size; + } return response.alloc_size; } @@ -673,9 +918,44 @@ static void ggml_backend_rpc_free(ggml_backend_t backend) { delete backend; } +static void ggml_backend_rpc_set_tensor_async(ggml_backend_t backend, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { + ggml_backend_rpc_context * ctx = (ggml_backend_rpc_context *)backend->context; + rpc_tensor rpc_tensor = serialize_tensor(tensor); + if (size > HASH_THRESHOLD) { + auto request = std::make_shared(); + request->tensor = rpc_tensor; + request->offset = offset; + request->hash = fnv_hash((const uint8_t*)data, size); + rpc_msg_set_tensor_hash_rsp response; + // TODO: make this async + ctx->dispatcher->send(RPC_CMD_SET_TENSOR_HASH, request, sizeof(*request), &response, sizeof(response)); + if (response.result) { + // the server has the same data, no need to send it + return; + } + } + // input serialization format: | rpc_tensor | offset (8 bytes) | data (size bytes) + size_t input_size = sizeof(rpc_tensor) + sizeof(uint64_t) + size; + uint8_t * input = new uint8_t[input_size](); + memcpy(input, &rpc_tensor, sizeof(rpc_tensor)); + memcpy(input + sizeof(rpc_tensor), &offset, sizeof(offset)); + memcpy(input + sizeof(rpc_tensor) + sizeof(offset), data, size); + std::shared_ptr input_ptr(input, std::default_delete()); + ctx->dispatcher->send_async(RPC_CMD_SET_TENSOR, input_ptr, input_size); +} + +static void ggml_backend_rpc_get_tensor_async(ggml_backend_t backend, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { + ggml_backend_rpc_context * ctx = (ggml_backend_rpc_context *)backend->context; + auto request = std::make_shared(); + request->tensor = serialize_tensor(tensor); + request->offset = offset; + request->size = size; + ctx->dispatcher->send_async(RPC_CMD_GET_TENSOR, request, sizeof(*request), data, size); +} + static void ggml_backend_rpc_synchronize(ggml_backend_t backend) { - GGML_UNUSED(backend); - // this is no-op because we don't have any async operations + ggml_backend_rpc_context * rpc_ctx = (ggml_backend_rpc_context *)backend->context; + rpc_ctx->dispatcher->synchronize(); } static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, std::vector & tensors, std::unordered_set & visited) { @@ -698,7 +978,7 @@ static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, std::ve tensors.push_back(result); } -static void serialize_graph(uint32_t device, const ggml_cgraph * cgraph, std::vector & output) { +static uint8_t * serialize_graph(uint32_t device, const ggml_cgraph * cgraph, size_t * output_size) { uint32_t n_nodes = cgraph->n_nodes; std::vector tensors; std::unordered_set visited; @@ -708,9 +988,9 @@ static void serialize_graph(uint32_t device, const ggml_cgraph * cgraph, std::ve // serialization format: // | device (4 bytes) | n_nodes (4 bytes) | nodes (n_nodes * sizeof(uint64_t) | n_tensors (4 bytes) | tensors (n_tensors * sizeof(rpc_tensor)) | uint32_t n_tensors = tensors.size(); - int output_size = 2*sizeof(uint32_t) + n_nodes * sizeof(uint64_t) + sizeof(uint32_t) + n_tensors * sizeof(rpc_tensor); - output.resize(output_size, 0); - uint8_t * dest = output.data(); + *output_size = 2*sizeof(uint32_t) + n_nodes * sizeof(uint64_t) + sizeof(uint32_t) + n_tensors * sizeof(rpc_tensor); + uint8_t * output = new uint8_t[*output_size](); + uint8_t * dest = output; memcpy(dest, &device, sizeof(device)); dest += sizeof(device); memcpy(dest, &n_nodes, sizeof(n_nodes)); @@ -723,6 +1003,7 @@ static void serialize_graph(uint32_t device, const ggml_cgraph * cgraph, std::ve dest += sizeof(n_tensors); rpc_tensor * out_tensors = (rpc_tensor *)dest; memcpy(out_tensors, tensors.data(), n_tensors * sizeof(rpc_tensor)); + return output; } static enum ggml_status ggml_backend_rpc_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) { @@ -733,27 +1014,35 @@ static enum ggml_status ggml_backend_rpc_graph_compute(ggml_backend_t backend, g GGML_ASSERT(cgraph->n_nodes > 0); bool reuse = cgraph->uid != 0 && rpc_dev_ctx->last_graph_uid == cgraph->uid; if (reuse) { - rpc_msg_graph_recompute_req request; - request.device = rpc_ctx->device; - auto sock = get_socket(rpc_ctx->endpoint); - bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_RECOMPUTE, &request, sizeof(request)); - RPC_STATUS_ASSERT(status); + auto request = std::make_shared(); + request->device = rpc_ctx->device; + rpc_ctx->dispatcher->send_async(RPC_CMD_GRAPH_RECOMPUTE, request, sizeof(*request)); } else { rpc_dev_ctx->last_graph_uid = cgraph->uid; - std::vector input; - serialize_graph(rpc_ctx->device, cgraph, input); - auto sock = get_socket(rpc_ctx->endpoint); - bool status = send_rpc_cmd(sock, RPC_CMD_GRAPH_COMPUTE, input.data(), input.size()); - RPC_STATUS_ASSERT(status); + size_t input_size = 0; + uint8_t * input = serialize_graph(rpc_ctx->device, cgraph, &input_size); + std::shared_ptr input_ptr(input, std::default_delete()); + rpc_ctx->dispatcher->send_async(RPC_CMD_GRAPH_COMPUTE, input_ptr, input_size); } return GGML_STATUS_SUCCESS; } +static void ggml_backend_rpc_event_record(ggml_backend_t backend, ggml_backend_event_t event) { + ggml_backend_rpc_context * rpc_ctx = (ggml_backend_rpc_context *)backend->context; + rpc_ctx->dispatcher->event_record(event); +} + +static void ggml_backend_rpc_event_wait(ggml_backend_t backend, ggml_backend_event_t event) { + // this is noop for RPC as we have a single stream + GGML_UNUSED(backend); + GGML_UNUSED(event); +} + static ggml_backend_i ggml_backend_rpc_interface = { /* .get_name = */ ggml_backend_rpc_name, /* .free = */ ggml_backend_rpc_free, - /* .set_tensor_async = */ NULL, - /* .get_tensor_async = */ NULL, + /* .set_tensor_async = */ ggml_backend_rpc_set_tensor_async, + /* .get_tensor_async = */ ggml_backend_rpc_get_tensor_async, /* .set_tensor_2d_async = */ NULL, /* .get_tensor_2d_async = */ NULL, /* .cpy_tensor_async = */ NULL, @@ -763,8 +1052,8 @@ static ggml_backend_i ggml_backend_rpc_interface = { /* .graph_plan_update = */ NULL, /* .graph_plan_compute = */ NULL, /* .graph_compute = */ ggml_backend_rpc_graph_compute, - /* .event_record = */ NULL, - /* .event_wait = */ NULL, + /* .event_record = */ ggml_backend_rpc_event_record, + /* .event_wait = */ ggml_backend_rpc_event_wait, /* .graph_optimize = */ NULL, }; @@ -778,13 +1067,9 @@ ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const char * endpoint, u if (it != buft_map.end()) { return it->second; } - auto sock = get_socket(endpoint); - if (sock == nullptr) { - GGML_LOG_ERROR("Failed to connect to %s\n", endpoint); - return nullptr; - } - size_t alignment = get_alignment(sock, device); - size_t max_size = get_max_size(sock, device); + auto dispatcher = get_dispatcher(endpoint); + size_t alignment = get_alignment(dispatcher, device); + size_t max_size = get_max_size(dispatcher, device); ggml_backend_rpc_buffer_type_context * buft_ctx = new ggml_backend_rpc_buffer_type_context { /* .endpoint = */ endpoint, /* .device = */ device, @@ -804,10 +1089,11 @@ ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const char * endpoint, u ggml_backend_t ggml_backend_rpc_init(const char * endpoint, uint32_t device) { std::string dev_name = "RPC" + std::to_string(device) + "[" + std::string(endpoint) + "]"; + auto dispatcher = get_dispatcher(endpoint); ggml_backend_rpc_context * ctx = new ggml_backend_rpc_context { - /* .endpoint = */ endpoint, - /* .device = */ device, - /* .name = */ dev_name, + /* .dispatcher = */ dispatcher, + /* .device = */ device, + /* .name = */ dev_name, }; auto reg = ggml_backend_rpc_add_server(endpoint); ggml_backend_t backend = new ggml_backend { @@ -823,26 +1109,16 @@ bool ggml_backend_is_rpc(ggml_backend_t backend) { return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_rpc_guid()); } -static void get_device_memory(const std::shared_ptr & sock, uint32_t device, size_t * free, size_t * total) { - rpc_msg_get_device_memory_req request; - request.device = device; +void ggml_backend_rpc_get_device_memory(const char * endpoint, uint32_t device, size_t * free, size_t * total) { + auto dispatcher = get_dispatcher(endpoint); + auto request = std::make_shared(); + request->device = device; rpc_msg_get_device_memory_rsp response; - bool status = send_rpc_cmd(sock, RPC_CMD_GET_DEVICE_MEMORY, &request, sizeof(request), &response, sizeof(response)); - RPC_STATUS_ASSERT(status); + dispatcher->send(RPC_CMD_GET_DEVICE_MEMORY, request, sizeof(*request), &response, sizeof(response)); *free = response.free_mem; *total = response.total_mem; } -void ggml_backend_rpc_get_device_memory(const char * endpoint, uint32_t device, size_t * free, size_t * total) { - auto sock = get_socket(endpoint); - if (sock == nullptr) { - *free = 0; - *total = 0; - return; - } - get_device_memory(sock, device, free, total); -} - // RPC server-side implementation class rpc_server { @@ -1647,9 +1923,6 @@ static void rpc_serve_client(const std::vector & backends, const if (!server.free_buffer(request)) { return; } - if (!send_msg(sock, nullptr, 0)) { - return; - } break; } case RPC_CMD_BUFFER_CLEAR: { @@ -1660,9 +1933,6 @@ static void rpc_serve_client(const std::vector & backends, const if (!server.buffer_clear(request)) { return; } - if (!send_msg(sock, nullptr, 0)) { - return; - } break; } case RPC_CMD_MEMSET_TENSOR: { @@ -1673,9 +1943,6 @@ static void rpc_serve_client(const std::vector & backends, const if (!server.memset_tensor(request)) { return; } - if (!send_msg(sock, nullptr, 0)) { - return; - } break; } case RPC_CMD_SET_TENSOR: { @@ -1710,9 +1977,6 @@ static void rpc_serve_client(const std::vector & backends, const if (!server.init_tensor(request)) { return; } - if (!send_msg(sock, nullptr, 0)) { - return; - } break; } case RPC_CMD_GET_TENSOR: { @@ -1889,10 +2153,10 @@ static void ggml_backend_rpc_device_get_props(ggml_backend_dev_t dev, struct ggm props->type = ggml_backend_rpc_device_get_type(dev); ggml_backend_rpc_device_get_memory(dev, &props->memory_free, &props->memory_total); props->caps = { - /* .async = */ false, + /* .async = */ true, /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ false, - /* .events = */ false, + /* .events = */ true, /* .mmap_support = */ true, }; } @@ -1929,6 +2193,24 @@ static bool ggml_backend_rpc_device_supports_buft(ggml_backend_dev_t dev, ggml_b return buft_ctx->endpoint == dev_ctx->endpoint && buft_ctx->device == dev_ctx->device; } +static ggml_backend_event_t ggml_backend_rpc_device_event_new(ggml_backend_dev_t dev) { + ggml_backend_rpc_device_context * ctx = (ggml_backend_rpc_device_context *)dev->context; + auto dispatcher = get_dispatcher(ctx->endpoint); + return dispatcher->event_new(dev); +} + +static void ggml_backend_rpc_device_event_free(ggml_backend_dev_t dev, ggml_backend_event_t event) { + ggml_backend_rpc_device_context * ctx = (ggml_backend_rpc_device_context *)dev->context; + auto dispatcher = get_dispatcher(ctx->endpoint); + dispatcher->event_free(event); +} + +static void ggml_backend_rpc_device_event_synchronize(ggml_backend_dev_t dev, ggml_backend_event_t event) { + ggml_backend_rpc_device_context * ctx = (ggml_backend_rpc_device_context *)dev->context; + auto dispatcher = get_dispatcher(ctx->endpoint); + dispatcher->event_synchronize(event); +} + static const struct ggml_backend_device_i ggml_backend_rpc_device_i = { /* .get_name = */ ggml_backend_rpc_device_get_name, /* .get_description = */ ggml_backend_rpc_device_get_description, @@ -1942,9 +2224,9 @@ static const struct ggml_backend_device_i ggml_backend_rpc_device_i = { /* .supports_op = */ ggml_backend_rpc_device_supports_op, /* .supports_buft = */ ggml_backend_rpc_device_supports_buft, /* .offload_op = */ NULL, - /* .event_new = */ NULL, - /* .event_free = */ NULL, - /* .event_synchronize = */ NULL, + /* .event_new = */ ggml_backend_rpc_device_event_new, + /* .event_free = */ ggml_backend_rpc_device_event_free, + /* .event_synchronize = */ ggml_backend_rpc_device_event_synchronize, }; // backend reg interface @@ -2004,14 +2286,9 @@ ggml_backend_reg_t ggml_backend_rpc_reg(void) { } static uint32_t ggml_backend_rpc_get_device_count(const char * endpoint) { - auto sock = get_socket(endpoint); - if (sock == nullptr) { - GGML_LOG_ERROR("Failed to connect to %s\n", endpoint); - return 0; - } + auto dispatcher = get_dispatcher(endpoint); rpc_msg_device_count_rsp response; - bool status = send_rpc_cmd(sock, RPC_CMD_DEVICE_COUNT, nullptr, 0, &response, sizeof(response)); - RPC_STATUS_ASSERT(status); + dispatcher->send(RPC_CMD_DEVICE_COUNT, nullptr, 0, &response, sizeof(response)); return response.device_count; } diff --git a/ggml/src/ggml-sycl/fattn-common.hpp b/ggml/src/ggml-sycl/fattn-common.hpp index c6cc13cfb00..82813f7a99a 100644 --- a/ggml/src/ggml-sycl/fattn-common.hpp +++ b/ggml/src/ggml-sycl/fattn-common.hpp @@ -6,6 +6,7 @@ #include "convert.hpp" #include "vecdotq.hpp" #include "fattn-buffers.hpp" +#include "fattn.hpp" #include "ggml.h" @@ -926,6 +927,7 @@ void launch_fattn( ggml_sycl_fattn_alloc K_f16(fbuf.K); ggml_sycl_fattn_alloc V_f16(fbuf.V); + const ggml_sycl_fattn_extra extra = ggml_sycl_fattn_get_extra(dst); ggml_sycl_pool_alloc KV_max(pool); ggml_sycl_pool_alloc dst_tmp(pool); ggml_sycl_pool_alloc dst_tmp_meta(pool); @@ -944,10 +946,11 @@ void launch_fattn( const size_t bs = ggml_blck_size(K->type); const size_t ts = ggml_type_size(K->type); - K_f16.alloc(ggml_nelements(K)); + sycl::half * K_f16_ptr = extra.K_buffer_ptr ? (sycl::half *) extra.K_buffer_ptr + : K_f16.alloc(ggml_nelements(K)); if (ggml_is_contiguously_allocated(K)) { to_fp16_sycl_t to_fp16 = ggml_get_to_fp16_sycl(K->type, dst); - to_fp16(K_data, K_f16.ptr, ggml_nelements(K), main_stream); + to_fp16(K_data, K_f16_ptr, ggml_nelements(K), main_stream); nb11 = nb11 * bs * sizeof(sycl::half) / ts; nb12 = nb12 * bs * sizeof(sycl::half) / ts; @@ -958,13 +961,13 @@ void launch_fattn( const int64_t s01 = nb11 / ts; const int64_t s02 = nb12 / ts; const int64_t s03 = nb13 / ts; - to_fp16(K_data, K_f16.ptr, K->ne[0], K->ne[1], K->ne[2], K->ne[3], s01, s02, s03, main_stream); + to_fp16(K_data, K_f16_ptr, K->ne[0], K->ne[1], K->ne[2], K->ne[3], s01, s02, s03, main_stream); nb11 = K->ne[0] * sizeof(sycl::half); nb12 = K->ne[1] * nb11; nb13 = K->ne[2] * nb12; } - K_data = (char *) K_f16.ptr; + K_data = (char *) K_f16_ptr; } if (need_f16_V && V->type != GGML_TYPE_F16) { @@ -977,11 +980,12 @@ void launch_fattn( const size_t bs = ggml_blck_size(V->type); const size_t ts = ggml_type_size(V->type); - V_f16.alloc(ggml_nelements(V)); + sycl::half * V_f16_ptr = extra.V_buffer_ptr ? (sycl::half *) extra.V_buffer_ptr + : V_f16.alloc(ggml_nelements(V)); if (ggml_is_contiguously_allocated(V)) { to_fp16_sycl_t to_fp16 = ggml_get_to_fp16_sycl(V->type, dst); - to_fp16(V_data, V_f16.ptr, ggml_nelements(V), main_stream); - V_data = (char *) V_f16.ptr; + to_fp16(V_data, V_f16_ptr, ggml_nelements(V), main_stream); + V_data = (char *) V_f16_ptr; nb21 = nb21 * bs * sizeof(sycl::half) / ts; nb22 = nb22 * bs * sizeof(sycl::half) / ts; @@ -992,13 +996,13 @@ void launch_fattn( const int64_t s01 = nb21 / ts; const int64_t s02 = nb22 / ts; const int64_t s03 = nb23 / ts; - to_fp16(V_data, V_f16.ptr, V->ne[0], V->ne[1], V->ne[2], V->ne[3], s01, s02, s03, main_stream); + to_fp16(V_data, V_f16_ptr, V->ne[0], V->ne[1], V->ne[2], V->ne[3], s01, s02, s03, main_stream); nb21 = V->ne[0] * sizeof(sycl::half); nb22 = V->ne[1] * nb21; nb23 = V->ne[2] * nb22; } - V_data = (char *) V_f16.ptr; + V_data = (char *) V_f16_ptr; } } diff --git a/ggml/src/ggml-sycl/fattn-onednn.cpp b/ggml/src/ggml-sycl/fattn-onednn.cpp index a501295192f..4349363a3d3 100644 --- a/ggml/src/ggml-sycl/fattn-onednn.cpp +++ b/ggml/src/ggml-sycl/fattn-onednn.cpp @@ -1,3 +1,4 @@ +#include #include #include #include @@ -13,9 +14,21 @@ // set minimum query length to treat as prefill (32) #define GGML_SYCL_FA_ONEDNN_MIN_Q 32 -bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst) { +bool ggml_sycl_fattn_onednn_binds_kv(const ggml_tensor * K, const ggml_tensor * V) { + if (K->type != GGML_TYPE_F16 || V->type != GGML_TYPE_F16) { + return false; + } + auto bindable = [](const ggml_tensor * t) { + return t->nb[0] == sizeof(sycl::half) && t->nb[1] % sizeof(sycl::half) == 0 && + t->nb[2] % sizeof(sycl::half) == 0 && t->nb[3] % sizeof(sycl::half) == 0; + }; + return bindable(K) && bindable(V); +} + +bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst, bool use_shape_limit) { #if !GGML_SYCL_DNNL GGML_UNUSED(dst); + GGML_UNUSED(use_shape_limit); return false; #else if (!g_ggml_sycl_fa_onednn) { @@ -43,7 +56,7 @@ bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst) { if (!k_ok || !v_ok) { return false; } - if (Q->ne[1] < 32 || K->ne[1] < 1024) { + if (use_shape_limit && (Q->ne[1] < 32 || K->ne[1] < 1024)) { return false; } for (const ggml_tensor * t : {K, V}) { @@ -93,7 +106,7 @@ bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst) { return false; } // Prefill only. - if (Q->ne[1] < GGML_SYCL_FA_ONEDNN_MIN_Q) { + if (use_shape_limit && Q->ne[1] < GGML_SYCL_FA_ONEDNN_MIN_Q) { return false; } return true; @@ -150,7 +163,8 @@ struct sdpa_partition { // Build + compile the contiguous-input GQA SDPA graph (MatMul->Divide->Add->SoftMax->MatMul), f32 out. // Mirrors the hardware-verified scratch/onednn_sdpa_probe.cpp build_gqa (partitions=1, sdp_primitive_kernel_t). -static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int seq, int d) { +static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int seq, int d, + const std::array & k_str, const std::array & v_str) try { using ltype = logical_tensor::layout_type; using dt = logical_tensor::data_type; using ldims = logical_tensor::dims; @@ -158,11 +172,12 @@ static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int const int rep = H / Hkv; const ldims q_sz = {1, Hkv, rep, q, d}, kv_sz = {1, Hkv, 1, seq, d}, s_sz = {1, Hkv, rep, q, seq}, sc = {1, 1, 1, 1, 1}, msk = {1, 1, 1, q, seq}, o_sz = {1, Hkv, rep, q, d}; + const ldims k_st(k_str.begin(), k_str.end()), v_st(v_str.begin(), v_str.end()); int64_t id = 0; sdpa_partition E; auto query = logical_tensor(id++, t, q_sz, ltype::strided); - auto key = logical_tensor(id++, t, kv_sz, ltype::strided); + auto key = logical_tensor(id++, t, kv_sz, k_st); auto score = logical_tensor(id++, fi, s_sz, ltype::strided); auto bmm1 = op(id++, op::kind::MatMul, "bmm1"); bmm1.set_attr(op::attr::transpose_b, true); // key is [.., seq, d] @@ -184,7 +199,7 @@ static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int smax.set_attr(op::attr::mode, "inf_as_zero"); smax.add_inputs({masked}); smax.add_outputs({probs}); - auto value = logical_tensor(id++, t, kv_sz, ltype::strided); + auto value = logical_tensor(id++, t, kv_sz, v_st); // f16 output is REQUIRED to hit sdp_primitive_kernel_t (the systolic micro-kernel); an f32 output // falls to larger_partition_kernel_t which materializes N^2 (confirmed: scratch/onednn_sdpa_kernel_probe.cpp). // converted to the f32 ggml dst in the permute below. @@ -198,6 +213,7 @@ static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int auto parts = g.get_partitions(); if (parts.size() != 1 || !parts[0].is_supported()) { + GGML_LOG_WARN("%s: oneDNN did not fuse the SDPA graph; falling back to TILE kernel\n", __func__); return E; // ok stays false -> caller falls back to TILE } E.ins = parts[0].get_input_ports(); @@ -209,6 +225,12 @@ static sdpa_partition build_sdpa(const engine & eng, int H, int Hkv, int q, int E.ok = true; return E; } +catch (const std::exception & e) { + // compile() can reject a stride set the partitioner never inspects; memoise the failure so the + // fallback costs one build rather than one per call. + GGML_LOG_WARN("%s: oneDNN SDPA partition build failed (%s); falling back to TILE kernel\n", __func__, e.what()); + return {}; +} void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tensor * dst) try { const ggml_tensor * Q = dst->src[0]; @@ -230,27 +252,53 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso dnnl::engine eng = ctx.engine_dnnl(stream); dnnl::stream strm = ctx.stream_dnnl(stream); + const ggml_sycl_fattn_extra extra = ggml_sycl_fattn_get_extra(dst); + // Q: always f32 -- copy to dense f16. - ggml_sycl_pool_alloc Qf(ctx.pool(), (size_t) H * q * d); - cont_to_f16_sycl((const char *) Q->data, Qf.get(), d, q, H, mb, Q->nb[1], Q->nb[2], Q->nb[3], stream); + std::optional> Qf_pool; + sycl::half * Qf_ptr = (sycl::half *) extra.Q_buffer_ptr; + if (!Qf_ptr) { + Qf_pool.emplace(ctx.pool(), (size_t) H * q * d); + Qf_ptr = Qf_pool->get(); + } + cont_to_f16_sycl((const char *) Q->data, Qf_ptr, d, q, H, mb, Q->nb[1], Q->nb[2], Q->nb[3], stream); - // K/V: use pool-alloc for both F16 and dequant paths. + // K/V: bind the f16 cache in place. llama.cpp permutes it to [token][head][dim], so its head + // plane is strided rather than dense, which is what an explicit stride vector expresses. + // Quantized and f32 KV still stage a dense copy -- the layout the k_str/v_str defaults describe. sycl::half * K_ptr = nullptr; sycl::half * V_ptr = nullptr; + std::array k_str{ Hkv * seq * d, seq * d, seq * d, d, 1 }; + std::array v_str = k_str; std::optional> Kf_pool; std::optional> Vf_pool; + // Helper: hand out reserved space, or fall back to the pool. + auto stage_k = [&](size_t n) { if (extra.K_buffer_ptr) { return (sycl::half *) extra.K_buffer_ptr; } + Kf_pool.emplace(ctx.pool(), n); return Kf_pool->get(); }; + auto stage_v = [&](size_t n) { if (extra.V_buffer_ptr) { return (sycl::half *) extra.V_buffer_ptr; } + Vf_pool.emplace(ctx.pool(), n); return Vf_pool->get(); }; + + auto elem_strides = [](const ggml_tensor * t) { + const int64_t s1 = (int64_t) (t->nb[1] / t->nb[0]); + const int64_t s2 = (int64_t) (t->nb[2] / t->nb[0]); + const int64_t s3 = (int64_t) (t->nb[3] / t->nb[0]); + // dims are {mb=1, Hkv, rep=1, seq, d}; the size-1 dims at 0 and 2 never advance an address. + return std::array{ s3, s2, s2, s1, 1 }; + }; - if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16) { - Kf_pool.emplace(ctx.pool(), (size_t) Hkv * seq * d); - Vf_pool.emplace(ctx.pool(), (size_t) Hkv * seq * d); - cont_to_f16_sycl((const char *) K->data, Kf_pool->get(), d, seq, Hkv, mb, K->nb[1], K->nb[2], K->nb[3], stream); - cont_to_f16_sycl((const char *) V->data, Vf_pool->get(), d, seq, Hkv, mb, V->nb[1], V->nb[2], V->nb[3], stream); - K_ptr = Kf_pool->get(); - V_ptr = Vf_pool->get(); + if (ggml_sycl_fattn_onednn_binds_kv(K, V)) { + K_ptr = (sycl::half *) K->data; + V_ptr = (sycl::half *) V->data; + k_str = elem_strides(K); + v_str = elem_strides(V); + } else if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16) { + K_ptr = stage_k((size_t) Hkv * seq * d); + V_ptr = stage_v((size_t) Hkv * seq * d); + cont_to_f16_sycl((const char *) K->data, K_ptr, d, seq, Hkv, mb, K->nb[1], K->nb[2], K->nb[3], stream); + cont_to_f16_sycl((const char *) V->data, V_ptr, d, seq, Hkv, mb, V->nb[1], V->nb[2], V->nb[3], stream); } else if (ggml_is_quantized(K->type)) { // Quantized K/V: dequant to dense F16 using pool, same lifetime as F16 path. - Kf_pool.emplace(ctx.pool(), ggml_nelements(K)); - K_ptr = Kf_pool->get(); + K_ptr = stage_k((size_t) ggml_nelements(K)); { const char * K_data = (const char *)K->data; const bool k_non_dense = ((int64_t)K->ne[1] * K->nb[1] != K->nb[2]) && K->ne[2] > 1; @@ -284,8 +332,7 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso // data pointer), their logical values differ because the quantized // elements at different positions/offsets represent different K/V // data. Master's F16 path also never aliases K and V. - Vf_pool.emplace(ctx.pool(), ggml_nelements(V)); - V_ptr = Vf_pool->get(); + V_ptr = stage_v((size_t) ggml_nelements(V)); { const char * V_data = (const char *)V->data; const bool v_non_dense = ((int64_t)V->ne[1] * V->nb[1] != V->nb[2]) && V->ne[2] > 1; @@ -316,12 +363,10 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso } } else { // F32: strided copy to dense F16 via cont_to_f16_sycl. - Kf_pool.emplace(ctx.pool(), ggml_nelements(K)); - K_ptr = Kf_pool->get(); + K_ptr = stage_k((size_t) ggml_nelements(K)); cont_to_f16_sycl((const char *) K->data, K_ptr, K->ne[0], K->ne[1], K->ne[2], K->ne[3], K->nb[1], K->nb[2], K->nb[3], stream); - Vf_pool.emplace(ctx.pool(), ggml_nelements(V)); - V_ptr = Vf_pool->get(); + V_ptr = stage_v((size_t) ggml_nelements(V)); cont_to_f16_sycl((const char *) V->data, V_ptr, V->ne[0], V->ne[1], V->ne[2], V->ne[3], V->nb[1], V->nb[2], V->nb[3], stream); } @@ -335,28 +380,43 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso // instead -- the value is captured into the command, so no host memory has to outlive the // call, and the enqueue stays async. const sycl::half scale_h = (sycl::half) (1.0f / kq_scale); - ggml_sycl_pool_alloc scbuf(ctx.pool(), 1); - sycl::half * const scale_dev = scbuf.get(); + std::optional> scbuf; + sycl::half * scale_dev = (sycl::half *) extra.scale_buffer_ptr; + if (!scale_dev) { + scbuf.emplace(ctx.pool(), 1); + scale_dev = scbuf->get(); + } stream->single_task([=]() { *scale_dev = scale_h; }); - ggml_sycl_pool_alloc outf(ctx.pool(), (size_t) H * q * d); // f16 contiguous SDPA out [mb,H,q,d] + // f16 contiguous SDPA out [mb,H,q,d] + std::optional> outf_pool; + sycl::half * outf_ptr = (sycl::half *) extra.out_buffer_ptr; + if (!outf_ptr) { + outf_pool.emplace(ctx.pool(), (size_t) H * q * d); + outf_ptr = outf_pool->get(); + } - // compile once per (device, shape), reuse across layers/calls. + // compile once per (device, shape, KV strides), reuse across layers/calls. Stride 2 always + // repeats stride 1 and stride 4 is always 1, so the key covers every entry that can differ. static std::unordered_map cache; - char keyb[96]; - snprintf(keyb, sizeof(keyb), "%d:%lld:%lld:%lld:%lld:%lld", ggml_sycl_get_device(), - (long long) H, (long long) Hkv, (long long) q, (long long) seq, (long long) d); + char keyb[256]; + snprintf(keyb, sizeof(keyb), "%d:%lld:%lld:%lld:%lld:%lld:%lld:%lld:%lld:%lld:%lld:%lld", ggml_sycl_get_device(), + (long long) H, (long long) Hkv, (long long) q, (long long) seq, (long long) d, + (long long) k_str[0], (long long) k_str[1], (long long) k_str[3], + (long long) v_str[0], (long long) v_str[1], (long long) v_str[3]); auto it = cache.find(keyb); if (it == cache.end()) { - it = cache.emplace(keyb, build_sdpa(eng, (int) H, (int) Hkv, (int) q, (int) seq, (int) d)).first; + it = cache.emplace(keyb, build_sdpa(eng, (int) H, (int) Hkv, (int) q, (int) seq, (int) d, k_str, v_str)).first; } sdpa_partition & E = it->second; - // _supported() is authoritative: if it accepted this op the partition must build. - // A failure here is a gap in _supported() -- surface it, don't mask it with a fallback. - GGML_ASSERT(E.ok && "oneDNN SDPA partition failed to build for a _supported() shape"); + if (!E.ok) { + // oneDNN can decline a shape or a stride set that _supported() never sees; build_sdpa warns per key. + ggml_sycl_flash_attn_ext_tile(ctx, dst); + return; + } auto id2ptr = [&](size_t r) -> void * { - if (r == E.id_q) return Qf.get(); + if (r == E.id_q) return Qf_ptr; if (r == E.id_k) return K_ptr; if (r == E.id_v) return V_ptr; if (r == E.id_scale) return scale_dev; @@ -368,10 +428,10 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso for (auto & lt : E.ins) { ti.emplace_back(lt, eng, id2ptr(lt.get_id())); } - tensor to(E.out, eng, outf.get()); + tensor to(E.out, eng, outf_ptr); E.cp.execute(strm, ti, {to}); - permute_sdpa_out_sycl(outf.get(), (float *) dst->data, mb, H, q, d, stream); + permute_sdpa_out_sycl(outf_ptr, (float *) dst->data, mb, H, q, d, stream); // Single device needs no sync: the dnnl stream wraps this same in-order queue, so the SDPA // serializes with the staging kernels before it and the permute/pool reuse after it. The // garbage output formerly blamed on the missing sync here was the scale use-after-return diff --git a/ggml/src/ggml-sycl/fattn-onednn.hpp b/ggml/src/ggml-sycl/fattn-onednn.hpp index d3019e87688..9669d1bd27a 100644 --- a/ggml/src/ggml-sycl/fattn-onednn.hpp +++ b/ggml/src/ggml-sycl/fattn-onednn.hpp @@ -5,7 +5,11 @@ // Static-only check: fused-XMX oneDNN Graph SDPA path==flash-attn op // (f16 KV, no softcap/ALiBi, single stream, tuned head_dim, prefill-sized q.) -bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst); +bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst, bool use_shape_limit = true); + +// True when the oneDNN path binds an F16 KV cache in place instead of staging a dense copy of +// it. Depends only on the types and strides of K and V, so the answer holds for every call. +bool ggml_sycl_fattn_onednn_binds_kv(const ggml_tensor * K, const ggml_tensor * V); // Run flash attention through oneDNN's fused xmx SDPA // execute the cached SDPA partition, write the f32 dst. Falls back to the TILE kernel on any failure. diff --git a/ggml/src/ggml-sycl/fattn.cpp b/ggml/src/ggml-sycl/fattn.cpp index a85eb721f6a..b73e6d46ffa 100644 --- a/ggml/src/ggml-sycl/fattn.cpp +++ b/ggml/src/ggml-sycl/fattn.cpp @@ -104,7 +104,6 @@ enum best_fattn_kernel { static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const ggml_tensor * dst) { - GGML_UNUSED(device); #ifndef SYCL_FLASH_ATTN GGML_UNUSED(dst); return BEST_FATTN_KERNEL_NONE; @@ -263,6 +262,11 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const } } else { if (Q->ne[1] <= 2) { + // TILE is faster for quantized KV decode on Xe2 (BMG); keep VEC on untested archs + const gpu_arch arch = ggml_sycl_info().devices[device].hw_info.arch; + if (arch == gpu_arch::intel_gpu_bmg_g21 || arch == gpu_arch::intel_gpu_bmg_g31) { + return BEST_FATTN_KERNEL_TILE; + } return BEST_FATTN_KERNEL_VEC; } } @@ -374,3 +378,76 @@ void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst bool ggml_sycl_flash_attn_ext_supported(int device, const ggml_tensor * dst) { return ggml_sycl_get_best_fattn_kernel(device, dst) != BEST_FATTN_KERNEL_NONE; } + +static uintptr_t ggml_sycl_fattn_reserve_halves(ggml_sycl_fattn_extra & extra, size_t n_halves) { + if (n_halves == 0) { + return 0; + } + extra.end = GGML_PAD(extra.end, SYCL_BUFFER_ALIGNMENT); + const uintptr_t block = extra.end; + extra.end += n_halves * sizeof(sycl::half); + return block; +} + +ggml_sycl_fattn_extra ggml_sycl_fattn_get_extra(const ggml_tensor * dst) { + ggml_sycl_fattn_extra extra; + + extra.end = (uintptr_t) dst->data + ggml_nbytes(dst); + + if (dst->op != GGML_OP_FLASH_ATTN_EXT) { + return extra; + } + + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + if (!Q || !K || !V) { + return extra; + } + + const int64_t d = K->ne[0]; + const int64_t H = Q->ne[2]; + const int64_t q = Q->ne[1]; + + // calculate the worst-case memory consumption across all kernels + const bool onednn_supported = ggml_sycl_flash_attn_ext_onednn_supported(dst, /* use_shape_limit */ false); + + const bool tile_needs_K = K->type != GGML_TYPE_F16; + const bool tile_needs_V = V->type != GGML_TYPE_F16; + + 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)); + + size_t need_K = 0, need_V = 0, need_Q = 0, need_out = 0, need_scale = 0; + if (onednn_supported) { + need_Q = (size_t) H * q * d; + need_out = (size_t) H * q * d; + need_scale = 1; + // an f16 cache is bound in place, so it needs no staging copy + if (!ggml_sycl_fattn_onednn_binds_kv(K, V)) { + need_K = (size_t) ggml_nelements(K); + need_V = (size_t) ggml_nelements(V); + } + } + if (tile_needs_K) { + need_K = std::max(need_K, (size_t) ggml_nelements(K)); + } + if (tile_needs_V) { + need_V = std::max(need_V, (size_t) ggml_nelements(V)); + } + + extra.Q_buffer_ptr = ggml_sycl_fattn_reserve_halves(extra, need_Q); + extra.K_buffer_ptr = ggml_sycl_fattn_reserve_halves(extra, need_K); + extra.V_buffer_ptr = (V_is_K_view && !onednn_supported && need_V) + ? extra.K_buffer_ptr + : ggml_sycl_fattn_reserve_halves(extra, need_V); + extra.scale_buffer_ptr = ggml_sycl_fattn_reserve_halves(extra, need_scale); + extra.out_buffer_ptr = ggml_sycl_fattn_reserve_halves(extra, need_out); + + return extra; +} + +size_t ggml_sycl_flash_attn_ext_get_alloc_size(const ggml_tensor * dst) { + const ggml_sycl_fattn_extra extra = ggml_sycl_fattn_get_extra(dst); + return (size_t) (extra.end - (uintptr_t) dst->data); +} diff --git a/ggml/src/ggml-sycl/fattn.hpp b/ggml/src/ggml-sycl/fattn.hpp index c093970a3fe..f803aa2a804 100644 --- a/ggml/src/ggml-sycl/fattn.hpp +++ b/ggml/src/ggml-sycl/fattn.hpp @@ -19,6 +19,24 @@ void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst bool ggml_sycl_flash_attn_ext_supported(int device, const ggml_tensor * dst); +// Scratch that flash attention needs beyond the output tensor +struct ggml_sycl_fattn_extra { + uintptr_t K_buffer_ptr = 0; // F16 copy of the K cache + uintptr_t V_buffer_ptr = 0; // F16 copy of the V cache + uintptr_t Q_buffer_ptr = 0; // dense F16 copy of Q, oneDNN only + uintptr_t scale_buffer_ptr = 0; // the softmax scale as an F16 scalar, oneDNN only + uintptr_t out_buffer_ptr = 0; // F16 SDPA output before conversion to F32, oneDNN only + uintptr_t end = 0; // one past the last reserved byte; sizes the allocation +}; + +// ggml_sycl_fattn_get_extra() is the single source of truth for the layout: it both sizes +// the reservation and hands out the pointers, so the two cannot disagree. +// Each field is the address of one reserved block, or 0 if that block was not reserved, +// in which case the caller allocates from the scratch pool instead. +ggml_sycl_fattn_extra ggml_sycl_fattn_get_extra(const ggml_tensor * dst); + +size_t ggml_sycl_flash_attn_ext_get_alloc_size(const ggml_tensor * dst); + void ggml_sycl_flash_attn_ext_mkl(ggml_backend_sycl_context & ctx, ggml_tensor * dst); #endif // GGML_SYCL_FATTN_HPP diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 0573643d834..dc8a1744323 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -955,7 +955,10 @@ static size_t ggml_backend_sycl_buffer_type_get_max_size(ggml_backend_buffer_typ } static size_t ggml_backend_sycl_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { - size_t size = ggml_nbytes(tensor); + // Reserve the additional scratch so it's visible to the graph allocator + size_t size = tensor->op == GGML_OP_FLASH_ATTN_EXT + ? ggml_sycl_flash_attn_ext_get_alloc_size(tensor) + : ggml_nbytes(tensor); int64_t ne0 = tensor->ne[0]; if (ggml_is_quantized(tensor->type)) { diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index c1d86aaac5c..39b4cd35980 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -767,6 +767,21 @@ static constexpr std::initializer_list> rms_norm_mul_rope_vie { 4, 0, 3 }, // set_rows->src[0] == view }; +static constexpr std::array lightning_indexer_k_types = { + GGML_TYPE_F32, + GGML_TYPE_F16, + GGML_TYPE_BF16, + GGML_TYPE_Q8_0, + GGML_TYPE_Q5_1, + GGML_TYPE_Q5_0, + GGML_TYPE_Q4_1, + GGML_TYPE_Q4_0, + GGML_TYPE_IQ4_NL, +}; + +static bool ggml_vk_lightning_indexer_k_type_supported(ggml_type type) { + return std::find(lightning_indexer_k_types.begin(), lightning_indexer_k_types.end(), type) != lightning_indexer_k_types.end(); +} struct vk_device_struct { std::recursive_mutex mutex; @@ -1042,6 +1057,8 @@ struct vk_device_struct { vk_pipeline pipeline_argsort_large_f32[num_argsort_pipelines]; vk_pipeline pipeline_topk_f32[num_topk_pipelines]; vk_pipeline pipeline_sum_rows_f32; + vk_pipeline pipeline_cross_entropy_loss_f32, pipeline_cross_entropy_loss_f32_wg512; + vk_pipeline pipeline_cross_entropy_loss_back_f32, pipeline_cross_entropy_loss_back_f32_wg512; vk_pipeline pipeline_fwht_f32[4]; vk_pipeline pipeline_cumsum_f32; vk_pipeline pipeline_cumsum_small_f32; @@ -1066,6 +1083,7 @@ struct vk_device_struct { vk_pipeline pipeline_rwkv_wkv6_f32; vk_pipeline pipeline_rwkv_wkv7_f32; vk_pipeline pipeline_gated_linear_attn_f32; + vk_pipeline pipeline_lightning_indexer_f32[GGML_TYPE_COUNT]; // [size_idx][kda] where size_idx: 0=d16, 1=d32, 2=d64, 3=d128 vk_pipeline pipeline_gated_delta_net[4][2]; vk_pipeline pipeline_ssm_scan_f32_d128; @@ -1329,7 +1347,8 @@ struct vk_mat_mat_id_push_constants { uint32_t stride_a; uint32_t stride_b; uint32_t stride_d; uint32_t batch_stride_a; uint32_t batch_stride_b; uint32_t batch_stride_d; uint32_t nei0; uint32_t nei1; uint32_t nbi1; uint32_t ne11; - uint32_t padded_N; + uint32_t n_experts; + uint32_t hoist_row_ids; }; struct vk_mat_vec_id_push_constants { uint32_t ncols; @@ -1410,6 +1429,10 @@ struct vk_op_count_experts_push_constants { uint32_t nb00; uint32_t nb01; uint32_t a_offset; + uint32_t n_experts; + uint32_t hoist_row_ids; + uint32_t ne00mp; + uint32_t ne00L; }; struct vk_op_glu_push_constants { @@ -1588,6 +1611,10 @@ template <> void init_pushconst_fastdiv(vk_op_glu_push_constants &p) { init_fastdiv_values(p.ne20, p.ne2_0mp, p.ne2_0L); } +template <> void init_pushconst_fastdiv(vk_op_count_experts_push_constants &p) { + init_fastdiv_values(p.ne00, p.ne00mp, p.ne00L); +} + struct vk_op_binary_push_constants { uint32_t ne; uint32_t ne00; uint32_t ne01; uint32_t ne02; uint32_t ne03; uint32_t nb00; uint32_t nb01; uint32_t nb02; uint32_t nb03; @@ -1846,6 +1873,26 @@ struct vk_op_gated_linear_attn_push_constants { uint32_t H; float scale; }; +struct vk_op_lightning_indexer_push_constants { + uint32_t n_kv; + uint32_t n_heads; + uint32_t n_tokens; + uint32_t n_streams; + uint32_t n_masks; + uint32_t dispatch_x; + uint32_t q_nb1; + uint32_t q_nb2; + uint32_t q_nb3; + uint32_t k_nb2; + uint32_t k_nb3; + uint32_t w_nb1; + uint32_t w_nb3; + uint32_t m_nb1; + uint32_t m_nb3; + uint32_t d_nb1; + uint32_t d_nb3; +}; +static_assert(sizeof(vk_op_lightning_indexer_push_constants) <= 128); struct vk_op_gated_delta_net_push_constants { uint32_t H; uint32_t n_tokens; @@ -2355,9 +2402,8 @@ struct ggml_backend_vk_context { // Cache most recent tensor that was converted into prealloc_y, and what pipeline it used to convert. vk_pipeline_struct * prealloc_y_last_pipeline_used {}; const ggml_tensor * prealloc_y_last_tensor_used {}; - // True when prealloc_y holds the padded fp16 layout used by the coopmat2 B decode-vector callback. - // If false, then it's contiguous. - bool prealloc_y_last_decode_vector_staging {}; + // True when the K dimension in prealloc_y is padded. + bool prealloc_y_last_k_padded {}; // Track which nodes have been used since the last sync, and whether they were written to std::vector unsynced_nodes_written; @@ -3902,11 +3948,16 @@ static vk_fa_pipeline_state get_fa_pipeline_state(const vk_device& device, const return vk_fa_pipeline_state{hsk, hsv, params.block_rows, params.block_cols, params.d_split, params.row_split, params.shmem_staging, params.path, params.workgroup_size, subgroup_size, aligned, f32acc, flags, params.limit_occupancy_shmem, k_type, v_type}; } +// Bytes per buffer block for the FaBlockBytesK/V spec constants. F32 is fed as +// a vec4 "block" of 4 floats, everything else uses its ggml block size. +static uint32_t fa_block_bytes(ggml_type t) { + if (t == GGML_TYPE_F32) { + return 16u; + } + return (uint32_t) ggml_type_size(t); +} + static std::vector get_fa_spec_constants(const vk_fa_pipeline_state& state) { - const auto fa_block_bytes = [](ggml_type t) -> uint32_t { - if (t == GGML_TYPE_F32) return 16u; - return (uint32_t) ggml_type_size(t); - }; return { /* 0 WorkGroupSize */ state.workgroup_size, /* 1 Br */ state.Br, @@ -4169,10 +4220,16 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { const uint32_t subgroup_size_16 = std::max(device->subgroup_size, 16u); const uint32_t subgroup_size_32 = std::max(device->subgroup_size, 32u); + // clamp WARP for l_/m_ warptiles so WM <= BM (breaks on subgroupSize > 64) + const uint32_t mm_warp_8 = std::min(subgroup_size_8, 64u); + const uint32_t mm_warp_16 = std::min(subgroup_size_16, 64u); + const uint32_t mul_mat_subgroup_size = (device->vendor_id == VK_VENDOR_ID_INTEL && device->subgroup_size_control) ? device->subgroup_min_size : device->subgroup_size; const uint32_t mul_mat_subgroup_size_8 = std::max(mul_mat_subgroup_size, 8u); const uint32_t mul_mat_subgroup_size_16 = std::max(mul_mat_subgroup_size, 16u); const uint32_t mul_mat_subgroup_size_32 = std::max(mul_mat_subgroup_size, 32u); + const uint32_t mul_mat_mm_warp_8 = std::min(mul_mat_subgroup_size_8, 64u); + const uint32_t mul_mat_mm_warp_16 = std::min(mul_mat_subgroup_size_16, 64u); const bool subgroup_min_size_16 = (!device->subgroup_size_control && device->subgroup_size >= 16) || (device->subgroup_size_control && device->subgroup_max_size >= 16); @@ -4253,39 +4310,39 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { const uint32_t s_warptile_wm = device->subgroup_size == 8 ? 8 : 32; - l_warptile = { 128, 128, 128, 16, subgroup_size_8 * 2, 64, 2, tm_l, tn_l, tk_l, subgroup_size_8 }; - m_warptile = { 128, 64, 64, 16, subgroup_size_8, 32, 2, tm_m, tn_m, tk_m, subgroup_size_8 }; - s_warptile = { subgroup_size_32, 32, 32, 16, s_warptile_wm, 32, 2, tm_s, tn_s, tk_s, subgroup_size_8 }; + l_warptile = { 128, 128, 128, 16, mm_warp_8 * 2, 64, 2, tm_l, tn_l, tk_l, mm_warp_8 }; + m_warptile = { 128, 64, 64, 16, mm_warp_8, 32, 2, tm_m, tn_m, tk_m, mm_warp_8 }; + s_warptile = { subgroup_size_32, 32, 32, 16, s_warptile_wm, 32, 2, tm_s, tn_s, tk_s, subgroup_size_8 }; - l_warptile_mmq = { 128, 128, 128, 32, subgroup_size_8 * 2, 64, 2, tm_l, tn_l, tk_l, subgroup_size_8 }; - m_warptile_mmq = { 128, 64, 64, 32, subgroup_size_8, 32, 2, tm_m, tn_m, tk_m, subgroup_size_8 }; - s_warptile_mmq = { subgroup_size_32, 32, 32, 32, s_warptile_wm, 32, 2, tm_s, tn_s, tk_s, subgroup_size_8 }; + l_warptile_mmq = { 128, 128, 128, 32, mm_warp_8 * 2, 64, 2, tm_l, tn_l, tk_l, mm_warp_8 }; + m_warptile_mmq = { 128, 64, 64, 32, mm_warp_8, 32, 2, tm_m, tn_m, tk_m, mm_warp_8 }; + s_warptile_mmq = { subgroup_size_32, 32, 32, 32, s_warptile_wm, 32, 2, tm_s, tn_s, tk_s, subgroup_size_8 }; // Integer MMQ has a smaller shared memory profile, but heavier register use - l_warptile_mmq_int = { 128, 128, 128, 32, subgroup_size_8 * 2, 64, 2, 4, 4, 1, subgroup_size_8 }; - m_warptile_mmq_int = { 128, 64, 64, 32, subgroup_size_8, 32, 2, 2, 2, 1, subgroup_size_8 }; - s_warptile_mmq_int = { subgroup_size_32, 32, 32, 32, s_warptile_wm, 32, 2, 2, 1, 1, subgroup_size_8 }; + l_warptile_mmq_int = { 128, 128, 128, 32, mm_warp_8 * 2, 64, 2, 4, 4, 1, mm_warp_8 }; + m_warptile_mmq_int = { 128, 64, 64, 32, mm_warp_8, 32, 2, 2, 2, 1, mm_warp_8 }; + s_warptile_mmq_int = { subgroup_size_32, 32, 32, 32, s_warptile_wm, 32, 2, 2, 1, 1, subgroup_size_8 }; // K-quants use even more registers, mitigate by setting WMITER to 1 - l_warptile_mmq_int_k = { 128, 128, 128, 32, subgroup_size_8 * 2, 64, 1, 4, 4, 1, subgroup_size_8 }; - m_warptile_mmq_int_k = { 128, 64, 64, 32, subgroup_size_8, 32, 1, 2, 2, 1, subgroup_size_8 }; - s_warptile_mmq_int_k = { subgroup_size_32, 32, 32, 32, s_warptile_wm, 32, 1, 2, 1, 1, subgroup_size_8 }; + l_warptile_mmq_int_k = { 128, 128, 128, 32, mm_warp_8 * 2, 64, 1, 4, 4, 1, mm_warp_8 }; + m_warptile_mmq_int_k = { 128, 64, 64, 32, mm_warp_8, 32, 1, 2, 2, 1, mm_warp_8 }; + s_warptile_mmq_int_k = { subgroup_size_32, 32, 32, 32, s_warptile_wm, 32, 1, 2, 1, 1, subgroup_size_8 }; - l_warptile_id = { 128, 128, 128, 16, mul_mat_subgroup_size_16 * 2, 64, 2, tm_l, tn_l, tk_l, mul_mat_subgroup_size_16 }; - m_warptile_id = { 128, 64, 64, 16, mul_mat_subgroup_size_16, 32, 2, tm_m, tn_m, tk_m, mul_mat_subgroup_size_16 }; - s_warptile_id = { mul_mat_subgroup_size_16, 32, 32, 16, s_warptile_wm, 32, 2, tm_s, tn_s, tk_s, mul_mat_subgroup_size_16 }; + l_warptile_id = { 128, 128, 128, 16, mul_mat_mm_warp_16 * 2, 64, 2, tm_l, tn_l, tk_l, mul_mat_mm_warp_16 }; + m_warptile_id = { 128, 64, 64, 16, mul_mat_mm_warp_16, 32, 2, tm_m, tn_m, tk_m, mul_mat_mm_warp_16 }; + s_warptile_id = { mul_mat_subgroup_size_16, 32, 32, 16, s_warptile_wm, 32, 2, tm_s, tn_s, tk_s, mul_mat_subgroup_size_16 }; - l_warptile_mmqid = { 128, 128, 128, 32, mul_mat_subgroup_size_8 * 2, 64, 2, tm_l, tn_l, tk_l, mul_mat_subgroup_size_8 }; - m_warptile_mmqid = { 128, 64, 64, 32, mul_mat_subgroup_size_8, 32, 2, tm_m, tn_m, tk_m, mul_mat_subgroup_size_8 }; - s_warptile_mmqid = { mul_mat_subgroup_size_32, 32, 32, 32, s_warptile_wm, 32, 2, tm_s, tn_s, tk_s, mul_mat_subgroup_size_8 }; + l_warptile_mmqid = { 128, 128, 128, 32, mul_mat_mm_warp_8 * 2, 64, 2, tm_l, tn_l, tk_l, mul_mat_mm_warp_8 }; + m_warptile_mmqid = { 128, 64, 64, 32, mul_mat_mm_warp_8, 32, 2, tm_m, tn_m, tk_m, mul_mat_mm_warp_8 }; + s_warptile_mmqid = { mul_mat_subgroup_size_32, 32, 32, 32, s_warptile_wm, 32, 2, tm_s, tn_s, tk_s, mul_mat_subgroup_size_8 }; - l_warptile_mmqid_int = { 128, 128, 128, 32, mul_mat_subgroup_size_8 * 2, 64, 2, 4, 4, 1, mul_mat_subgroup_size_8 }; - m_warptile_mmqid_int = { 128, 64, 64, 32, mul_mat_subgroup_size_8, 32, 2, 2, 2, 1, mul_mat_subgroup_size_8 }; - s_warptile_mmqid_int = { mul_mat_subgroup_size_32, 32, 32, 32, s_warptile_wm, 32, 2, 2, 1, 1, mul_mat_subgroup_size_8 }; + l_warptile_mmqid_int = { 128, 128, 128, 32, mul_mat_mm_warp_8 * 2, 64, 2, 4, 4, 1, mul_mat_mm_warp_8 }; + m_warptile_mmqid_int = { 128, 64, 64, 32, mul_mat_mm_warp_8, 32, 2, 2, 2, 1, mul_mat_mm_warp_8 }; + s_warptile_mmqid_int = { mul_mat_subgroup_size_32, 32, 32, 32, s_warptile_wm, 32, 2, 2, 1, 1, mul_mat_subgroup_size_8 }; - l_warptile_mmqid_int_k = { 128, 128, 128, 32, mul_mat_subgroup_size_16 * 2, 64, 1, 4, 4, 1, mul_mat_subgroup_size_16 }; - m_warptile_mmqid_int_k = { 128, 64, 64, 32, mul_mat_subgroup_size_16, 32, 1, 2, 2, 1, mul_mat_subgroup_size_16 }; - s_warptile_mmqid_int_k = { mul_mat_subgroup_size_32, 32, 32, 32, s_warptile_wm, 32, 1, 2, 1, 1, mul_mat_subgroup_size_16 }; + l_warptile_mmqid_int_k = { 128, 128, 128, 32, mul_mat_mm_warp_16 * 2, 64, 1, 4, 4, 1, mul_mat_mm_warp_16 }; + m_warptile_mmqid_int_k = { 128, 64, 64, 32, mul_mat_mm_warp_16, 32, 1, 2, 2, 1, mul_mat_mm_warp_16 }; + s_warptile_mmqid_int_k = { mul_mat_subgroup_size_32, 32, 32, 32, s_warptile_wm, 32, 1, 2, 1, 1, mul_mat_subgroup_size_16 }; // chip specific tuning if ((device->architecture == AMD_GCN) && (device->driver_id != vk::DriverId::eAmdProprietary)) { @@ -4293,13 +4350,13 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { m_warptile_mmqid = m_warptile_mmqid_int = { 256, 64, 64, 32, 16, 16, 2, 2, 2, 1, 16 }; } else if (device->vendor_id == VK_VENDOR_ID_AMD && device->coopmat_support && device->driver_id != vk::DriverId::eAmdProprietary) { // This is intentionally using tx_m values, slight performance increase - l_warptile = { 256, 128, 128, 16, subgroup_size_8, 64, 2, tm_m, tn_m, tk_m, subgroup_size_8 }; - l_warptile_mmq = l_warptile_mmq_int = { 256, 128, 128, 32, subgroup_size_8, 64, 2, tm_m, tn_m, tk_m, subgroup_size_8 }; - l_warptile_mmq_int_k = { 256, 128, 128, 32, subgroup_size_16, 64, 1, 4, 2, 1, subgroup_size_16 }; + l_warptile = { 256, 128, 128, 16, mm_warp_8, 64, 2, tm_m, tn_m, tk_m, mm_warp_8 }; + l_warptile_mmq = l_warptile_mmq_int = { 256, 128, 128, 32, mm_warp_8, 64, 2, tm_m, tn_m, tk_m, mm_warp_8 }; + l_warptile_mmq_int_k = { 256, 128, 128, 32, mm_warp_16, 64, 1, 4, 2, 1, mm_warp_16 }; } else if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support) { // Xe2/Xe3 with coopmat enabled - warptile performance tuning - l_warptile = { 512, 128, 128, 16, subgroup_size_8, 32, 2, tm_m, tn_m, tk_m, subgroup_size_8 }; - l_warptile_mmq = { 512, 128, 128, 32, subgroup_size_8, 32, 2, tm_m, tn_m, tk_m, subgroup_size_8 }; + l_warptile = { 512, 128, 128, 16, mm_warp_8, 32, 2, tm_m, tn_m, tk_m, mm_warp_8 }; + l_warptile_mmq = { 512, 128, 128, 32, mm_warp_8, 32, 2, tm_m, tn_m, tk_m, mm_warp_8 }; } l_mmq_wg_denoms = l_wg_denoms = {128, 128, 1 }; @@ -5172,8 +5229,8 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { const uint32_t s_warptile_wm = device->subgroup_size == 8 ? 8 : 32; // use scalar tile sizes - l_warptile = { 128, 128, 128, 16, subgroup_size_8 * 2, 64, 2, 4, 4, 1, subgroup_size_8 }; - m_warptile = { 128, 64, 64, 16, subgroup_size_8, 32, 2, 4, 2, 1, subgroup_size_8 }; + l_warptile = { 128, 128, 128, 16, mm_warp_8 * 2, 64, 2, 4, 4, 1, mm_warp_8 }; + m_warptile = { 128, 64, 64, 16, mm_warp_8, 32, 2, 4, 2, 1, mm_warp_8 }; s_warptile = { subgroup_size_32, 32, 32, 16, s_warptile_wm, 32, 2, 2, 2, 1, subgroup_size_8 }; l_wg_denoms = {128, 128, 1 }; @@ -5758,6 +5815,10 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_argmax_f32, "argmax_f32", argmax_f32_len, argmax_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); ggml_vk_create_pipeline(device, device->pipeline_sum_rows_f32, "sum_rows_f32", sum_rows_f32_len, sum_rows_f32_data, "main", 2, sizeof(vk_op_sum_rows_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); + ggml_vk_create_pipeline(device, device->pipeline_cross_entropy_loss_f32, "cross_entropy_loss_f32", cross_entropy_loss_f32_len, cross_entropy_loss_f32_data, "main", 3, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); + ggml_vk_create_pipeline(device, device->pipeline_cross_entropy_loss_f32_wg512, "cross_entropy_loss_f32_wg512", cross_entropy_loss_f32_len, cross_entropy_loss_f32_data, "main", 3, sizeof(vk_op_push_constants), {1, 1, 1}, { 512 }, 1); + ggml_vk_create_pipeline(device, device->pipeline_cross_entropy_loss_back_f32, "cross_entropy_loss_back_f32", cross_entropy_loss_back_f32_len, cross_entropy_loss_back_f32_data, "main", 4, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); + ggml_vk_create_pipeline(device, device->pipeline_cross_entropy_loss_back_f32_wg512, "cross_entropy_loss_back_f32_wg512", cross_entropy_loss_back_f32_len, cross_entropy_loss_back_f32_data, "main", 4, sizeof(vk_op_push_constants), {1, 1, 1}, { 512 }, 1); // Intel Windows driver in range [32.0.101.8509, 32.0.101.8860) will crash when using fwht kernels so we gate that here const bool can_use_fwht = device->driver_id != vk::DriverId::eIntelProprietaryWindows || !ggml_vk_intel_windows_driver_in_range(device->properties.driverVersion, 101, 8509, 101, 8860); @@ -5786,7 +5847,11 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_count_equal_i32, "count_equal_i32", count_equal_i32_len, count_equal_i32_data, "main", 3, sizeof(vk_op_push_constants), {512, 1, 1}, { device->subgroup_size }, 1); - ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_len, count_experts_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true); + if (device->subgroup_arithmetic && device->subgroup_require_full_support) { + ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_subgroup_len, count_experts_subgroup_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true, true); + } else { + ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_len, count_experts_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true); + } for (auto &s : device->pipeline_solve_tri_f32) { const vk_solve_tri_pipeline_state &state = s.first; @@ -5835,6 +5900,17 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_gated_linear_attn_f32, "gated_linear_attn_f32", gated_linear_attn_f32_len, gated_linear_attn_f32_data, "main", 6, sizeof(vk_op_gated_linear_attn_push_constants), {1, 1, 1}, {}, 1); + { + const bool li_subgroup = device->subgroup_arithmetic && device->subgroup_require_full_support; + const size_t li_len = li_subgroup ? lightning_indexer_subgroup_f32_len : lightning_indexer_f32_len; + const void * li_data = li_subgroup ? (const void *)lightning_indexer_subgroup_f32_data : (const void *)lightning_indexer_f32_data; + + for (ggml_type k_type : lightning_indexer_k_types) { + const std::string name = "lightning_indexer_" + std::string(ggml_type_name(k_type)) + "_k_f32"; + ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_f32[k_type], name.c_str(), li_len, li_data, "main", 5, sizeof(vk_op_lightning_indexer_push_constants), {1, 1, 1}, {(uint32_t)k_type, fa_block_bytes(k_type), device->subgroup_size}, 1, true, li_subgroup); + } + } + { const uint32_t gdn_sizes[] = {16, 32, 64, 128}; const char * gdn_names[][2] = { @@ -8906,13 +8982,13 @@ static void ggml_vk_matmul_id( uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d, uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d, uint32_t n_as, uint32_t nei0, uint32_t nei1, uint32_t nbi1, uint32_t ne11, - uint32_t padded_n) { + bool hoist_row_ids) { VK_LOG_DEBUG("ggml_vk_matmul_id(a: (" << a.buffer->buffer << ", " << a.offset << ", " << a.size << "), b: (" << b.buffer->buffer << ", " << b.offset << ", " << b.size << "), d: (" << d.buffer->buffer << ", " << d.offset << ", " << d.size << "), ids: (" << ids.buffer->buffer << ", " << ids.offset << ", " << ids.size << "), expert_count: (" << expert_count_buf.buffer->buffer << ", " << expert_count_buf.offset << ", " << expert_count_buf.size << "), " << "m: " << m << ", n: " << n << ", k: " << k << ", stride_a: " << stride_a << ", stride_b: " << stride_b << ", stride_d: " << stride_d << ", " << "batch_stride_a: " << batch_stride_a << ", batch_stride_b: " << batch_stride_b << ", batch_stride_d: " << batch_stride_d << ", " << "n_as: " << n_as << ", nei0: " << nei0 << ", nei1: " << nei1 << ", nbi1: " << nbi1 << ", ne11: " << ne11 << ")"); const vk_mat_mat_id_push_constants pc = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d, - nei0, nei1, nbi1, ne11, padded_n }; + nei0, nei1, nbi1, ne11, n_as, uint32_t(hoist_row_ids) }; ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d, ids, expert_count_buf }, pc, { m, nei1, n_as }); } @@ -9377,27 +9453,27 @@ static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& sub if (y_non_contig) { if (ctx->prealloc_y_last_pipeline_used != to_fp16_vk_1.get() || ctx->prealloc_y_last_tensor_used != src1 || - ctx->prealloc_y_last_decode_vector_staging) { + ctx->prealloc_y_last_k_padded) { if (ctx->prealloc_y_need_sync) { ggml_vk_sync_buffers(ctx, subctx); } ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, ggml_vk_subbuffer(ctx, d_Qy, qy_buf_offset), ggml_vk_subbuffer(ctx, d_Y, 0)); ctx->prealloc_y_last_pipeline_used = to_fp16_vk_1.get(); ctx->prealloc_y_last_tensor_used = src1; - ctx->prealloc_y_last_decode_vector_staging = false; + ctx->prealloc_y_last_k_padded = false; } } if (quantize_y) { if (ctx->prealloc_y_last_pipeline_used != to_q8_1.get() || ctx->prealloc_y_last_tensor_used != src1 || - ctx->prealloc_y_last_decode_vector_staging) { + ctx->prealloc_y_last_k_padded) { if (ctx->prealloc_y_need_sync) { ggml_vk_sync_buffers(ctx, subctx); } ggml_vk_quantize_q8_1(ctx, subctx, ggml_vk_subbuffer(ctx, d_Qy, qy_buf_offset), ggml_vk_subbuffer(ctx, d_Y, 0), y_ne); ctx->prealloc_y_last_pipeline_used = to_q8_1.get(); ctx->prealloc_y_last_tensor_used = src1; - ctx->prealloc_y_last_decode_vector_staging = false; + ctx->prealloc_y_last_k_padded = false; } } @@ -9656,27 +9732,27 @@ static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context& GGML_ASSERT(y_sz == ggml_type_size(src1->type) * y_ne); if (ctx->prealloc_y_last_pipeline_used != to_fp16_vk_1.get() || ctx->prealloc_y_last_tensor_used != src1 || - ctx->prealloc_y_last_decode_vector_staging) { + ctx->prealloc_y_last_k_padded) { if (ctx->prealloc_y_need_sync) { ggml_vk_sync_buffers(ctx, subctx); } ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, d_Qy, d_Y); ctx->prealloc_y_last_pipeline_used = to_fp16_vk_1.get(); ctx->prealloc_y_last_tensor_used = src1; - ctx->prealloc_y_last_decode_vector_staging = false; + ctx->prealloc_y_last_k_padded = false; } } if (quantize_y) { if (ctx->prealloc_y_last_pipeline_used != to_q8_1.get() || ctx->prealloc_y_last_tensor_used != src1 || - ctx->prealloc_y_last_decode_vector_staging) { + ctx->prealloc_y_last_k_padded) { if (ctx->prealloc_y_need_sync) { ggml_vk_sync_buffers(ctx, subctx); } ggml_vk_quantize_q8_1(ctx, subctx, d_Qy, d_Y, y_ne); ctx->prealloc_y_last_pipeline_used = to_q8_1.get(); ctx->prealloc_y_last_tensor_used = src1; - ctx->prealloc_y_last_decode_vector_staging = false; + ctx->prealloc_y_last_k_padded = false; } } @@ -10098,6 +10174,12 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& // const uint64_t ne23 = dst->ne[3]; const uint64_t n_as = ne02; + // n_as counts, n_as offsets, one total, then one packed row id per (expert, token). + // Hoisting requires 16-bit indices for the packing and a table that fits one binding. + const uint64_t hoisted_row_id_words = 2 * n_as + 1 + nei0 * nei1; + const bool hoist_row_ids = n_as <= 256 && nei0 <= 0xffff && nei1 <= 0xffff && + hoisted_row_id_words * sizeof(uint32_t) <= + ctx->device->properties.limits.maxStorageBufferRange; ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; @@ -10150,8 +10232,6 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& (src0->type == GGML_TYPE_BF16 && src1->type != GGML_TYPE_BF16) || !ggml_vk_dim01_contiguous(src1); - const uint32_t y_staged_row_stride = y_decode_vector_staging ? (uint32_t)ggml_vk_align_size(ne10, 4) : (uint32_t)ne10; - const bool y_f32_kernel = src1->type == GGML_TYPE_F32 && !y_non_contig; bool quantize_y = ctx->device->integer_dot_product && src1->type == GGML_TYPE_F32 && ggml_is_contiguous(src1) && !y_non_contig && (ne11 * ne10) % 4 == 0; @@ -10166,19 +10246,25 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& } const bool qx_needs_dequant = mmp == nullptr || x_non_contig; - const bool qy_needs_dequant = !quantize_y && ((src1->type != f16_type && !y_f32_kernel) || y_non_contig); + bool qy_needs_dequant = !quantize_y && ((src1->type != f16_type && !y_f32_kernel) || y_non_contig); if (qx_needs_dequant) { // Fall back to dequant + f16 mulmat mmp = ggml_vk_get_mul_mat_mat_id_pipeline(ctx, f16_type, y_f32_kernel ? GGML_TYPE_F32 : f16_type, (ggml_prec)dst->op_params[0]); } - // Not implemented - GGML_ASSERT(y_non_contig || !qy_needs_dequant); // NOLINT - const ggml_type effective_src1_type = quantize_y ? GGML_TYPE_Q8_1 : (y_f32_kernel ? GGML_TYPE_F32 : src1->type); const uint32_t kpad = quantize_y ? 0 : ggml_vk_align_size(ne10, ggml_vk_guess_matmul_id_pipeline_align(ctx, mmp, ne01, nei1, qx_needs_dequant ? f16_type : src0->type, effective_src1_type)); + // Coopmat2 MUL_MAT_ID BK specialization constants in ggml_vk_load_shaders are at most 64. + const uint32_t y_staged_row_stride = ctx->device->coopmat2 && !quantize_y ? ggml_vk_align_size(ne10, 64) : ne10; + const bool y_needs_k_padding = ne10 != y_staged_row_stride; + const bool y_needs_reformat = y_non_contig || y_needs_k_padding; + qy_needs_dequant = qy_needs_dequant || y_needs_k_padding; + + // Not implemented + GGML_ASSERT(y_needs_reformat || !qy_needs_dequant); // NOLINT + const bool aligned = !quantize_y && ne10 == kpad && ne01 > 8 && nei1 > 8; vk_pipeline pipeline = ggml_vk_guess_matmul_id_pipeline(ctx, mmp, ne01, nei1, aligned, qx_needs_dequant ? f16_type : src0->type, effective_src1_type); @@ -10186,10 +10272,8 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& if (ggml_nbytes(src0) > ctx->device->properties.limits.maxStorageBufferRange) { pipeline = ggml_vk_get_64b_indexing_pipeline(ctx, pipeline); } - // Reserve extra storage in the N dimension for the Y matrix, so we can avoid bounds-checking - uint32_t padded_n = qy_needs_dequant ? ROUNDUP_POW2(ne11, pipeline->wg_denoms[1]) :ne11; const uint64_t x_ne = ggml_nelements(src0); - const uint64_t y_ne = (uint64_t)y_staged_row_stride * padded_n * ne12 * ne13; + const uint64_t y_ne = (uint64_t)y_staged_row_stride * ne11 * ne12 * ne13; const uint64_t d_ne = ggml_nelements(dst); const uint64_t qx_sz = ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type); @@ -10208,7 +10292,7 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& y_staged_dst.type = f16_type; y_staged_dst.nb[0] = ggml_type_size(f16_type); y_staged_dst.nb[1] = y_staged_dst.nb[0] * y_staged_row_stride; - y_staged_dst.nb[2] = y_staged_dst.nb[1] * padded_n; + y_staged_dst.nb[2] = y_staged_dst.nb[1] * ne11; y_staged_dst.nb[3] = y_staged_dst.nb[2] * y_staged_dst.ne[2]; return y_staged_dst; }; @@ -10218,10 +10302,10 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& } else { to_fp16_vk_0 = ggml_vk_get_to_fp16(ctx, src0->type); } - if (y_non_contig) { + if (y_needs_reformat) { ggml_tensor y_staged_dst; const ggml_tensor * y_staged_dst_ptr = nullptr; - if (y_decode_vector_staging) { + if (y_needs_k_padding) { y_staged_dst = make_y_staged_dst(); y_staged_dst_ptr = &y_staged_dst; } @@ -10238,7 +10322,8 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& } vk_pipeline count_experts = ctx->device->pipeline_count_experts; - uint32_t expert_count_size = sizeof(uint32_t) * n_as; + const size_t expert_data_size = sizeof(uint32_t) * + (hoist_row_ids ? hoisted_row_id_words : n_as); { if ( @@ -10254,8 +10339,8 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& ctx->prealloc_size_y = y_sz; ggml_vk_preallocate_buffers(ctx, subctx); } - if (ctx->prealloc_size_split_k < expert_count_size) { - ctx->prealloc_size_split_k = expert_count_size; + if (ctx->prealloc_size_split_k < expert_data_size) { + ctx->prealloc_size_split_k = expert_data_size; ggml_vk_preallocate_buffers(ctx, subctx); } @@ -10321,18 +10406,23 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& } } // Count how many times each expert is used - vk_subbuffer expert_count_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0); + vk_subbuffer expert_count_buf = { ctx->prealloc_split_k, 0, expert_data_size }; if (ctx->prealloc_split_k_need_sync) { ggml_vk_sync_buffers(ctx, subctx); } { - const std::vector pc = { (uint32_t)nei0, + vk_op_count_experts_push_constants pc = { (uint32_t)nei0, (uint32_t)nei1, (uint32_t)(nbi0 / ggml_type_size(ids->type)), (uint32_t)(nbi1 / ggml_type_size(ids->type)), - (uint32_t)(get_misalign_bytes(ctx, ids) / ggml_type_size(ids->type)) }; + (uint32_t)(get_misalign_bytes(ctx, ids) / ggml_type_size(ids->type)), + (uint32_t)n_as, + uint32_t(hoist_row_ids), + 0, 0 }; + init_pushconst_fastdiv(pc); ggml_vk_dispatch_pipeline(ctx, subctx, count_experts, - { vk_subbuffer{ d_ids, ids_buf_offset, ids_sz }, expert_count_buf }, pc, { (uint32_t)n_as, 1, 1}); + { vk_subbuffer{ d_ids, ids_buf_offset, ids_sz }, expert_count_buf }, pc, + { hoist_row_ids ? 1u : (uint32_t)n_as, 1, 1}); } if (x_non_contig) { @@ -10342,14 +10432,18 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& ggml_vk_dispatch_pipeline(ctx, subctx, to_fp16_vk_0, { vk_subbuffer{ d_Qx, qx_buf_offset, qx_sz }, vk_subbuffer{ d_X, 0, x_sz } }, pc, { (uint32_t)x_ne, 1, 1}); } - if (y_non_contig) { + if (y_needs_reformat) { if (ctx->prealloc_y_last_pipeline_used != to_fp16_vk_1.get() || ctx->prealloc_y_last_tensor_used != src1 || - ctx->prealloc_y_last_decode_vector_staging != y_decode_vector_staging) { + ctx->prealloc_y_last_k_padded != y_needs_k_padding) { if (ctx->prealloc_y_need_sync) { ggml_vk_sync_buffers(ctx, subctx); } - if (y_decode_vector_staging) { + if (y_needs_k_padding) { + GGML_ASSERT(y_sz % 4 == 0); + // Zero B padding because clamping only A can produce 0 * Inf or NaN. + subctx->s->buffer->buf.fillBuffer(d_Y->buffer, 0, y_sz, 0); + ggml_vk_sync_buffers(ctx, subctx); const ggml_tensor y_staged_dst = make_y_staged_dst(); const uint32_t y_staged_dst_type_size = ggml_type_size(y_staged_dst.type); ggml_vk_cpy_to_strided( @@ -10364,27 +10458,27 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& } ctx->prealloc_y_last_pipeline_used = to_fp16_vk_1.get(); ctx->prealloc_y_last_tensor_used = src1; - ctx->prealloc_y_last_decode_vector_staging = y_decode_vector_staging; + ctx->prealloc_y_last_k_padded = y_needs_k_padding; } } if (quantize_y) { if (ctx->prealloc_y_last_pipeline_used != to_q8_1.get() || ctx->prealloc_y_last_tensor_used != src1 || - ctx->prealloc_y_last_decode_vector_staging) { + ctx->prealloc_y_last_k_padded) { if (ctx->prealloc_y_need_sync) { ggml_vk_sync_buffers(ctx, subctx); } ggml_vk_quantize_q8_1(ctx, subctx, ggml_vk_subbuffer(ctx, d_Qy, qy_buf_offset), ggml_vk_subbuffer(ctx, d_Y, 0), y_ne); ctx->prealloc_y_last_pipeline_used = to_q8_1.get(); ctx->prealloc_y_last_tensor_used = src1; - ctx->prealloc_y_last_decode_vector_staging = false; + ctx->prealloc_y_last_k_padded = false; } } ggml_vk_sync_buffers(ctx, subctx); uint32_t stride_batch_x = ne00*ne01; - uint32_t stride_b_y = y_decode_vector_staging ? y_staged_row_stride : ne10; - uint32_t stride_batch_y = y_decode_vector_staging ? y_staged_row_stride * padded_n : ne10*ne11; + uint32_t stride_b_y = y_needs_k_padding ? y_staged_row_stride : ne10; + uint32_t stride_batch_y = y_needs_k_padding ? y_staged_row_stride * ne11 : ne10*ne11; if (!ggml_vk_dim01_contiguous(src0) && !qx_needs_dequant) { stride_batch_x = src0->nb[0] / ggml_type_size(src0->type); @@ -10401,13 +10495,13 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& { d_D, d_buf_offset, d_sz }, { d_ids, ids_buf_offset, ids_sz }, expert_count_buf, ne01, ne21, ne10, ne10, stride_b_y, ne01, stride_batch_x, stride_batch_y, ne20*ne21, - n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11, padded_n + n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11, hoist_row_ids ); // NOLINT if (x_non_contig || qx_needs_dequant) { ctx->prealloc_x_need_sync = true; } - if (y_non_contig || quantize_y) { + if (y_needs_reformat || quantize_y) { ctx->prealloc_y_need_sync = true; } ctx->prealloc_split_k_need_sync = true; @@ -10558,27 +10652,27 @@ static void ggml_vk_mul_mat_vec_id_q_f16(ggml_backend_vk_context * ctx, vk_conte GGML_ASSERT(y_sz == ggml_type_size(src1->type) * y_ne); if (ctx->prealloc_y_last_pipeline_used != to_fp16_vk_1.get() || ctx->prealloc_y_last_tensor_used != src1 || - ctx->prealloc_y_last_decode_vector_staging) { + ctx->prealloc_y_last_k_padded) { if (ctx->prealloc_y_need_sync) { ggml_vk_sync_buffers(ctx, subctx); } ggml_vk_cpy_to_contiguous(ctx, subctx, to_fp16_vk_1, src1, d_Qy, d_Y); ctx->prealloc_y_last_pipeline_used = to_fp16_vk_1.get(); ctx->prealloc_y_last_tensor_used = src1; - ctx->prealloc_y_last_decode_vector_staging = false; + ctx->prealloc_y_last_k_padded = false; } } if (quantize_y) { if (ctx->prealloc_y_last_pipeline_used != to_q8_1.get() || ctx->prealloc_y_last_tensor_used != src1 || - ctx->prealloc_y_last_decode_vector_staging) { + ctx->prealloc_y_last_k_padded) { if (ctx->prealloc_y_need_sync) { ggml_vk_sync_buffers(ctx, subctx); } ggml_vk_quantize_q8_1(ctx, subctx, d_Qy, d_Y, y_ne); ctx->prealloc_y_last_pipeline_used = to_q8_1.get(); ctx->prealloc_y_last_tensor_used = src1; - ctx->prealloc_y_last_decode_vector_staging = false; + ctx->prealloc_y_last_k_padded = false; } } @@ -11577,6 +11671,17 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const return ctx->device->pipeline_sum_rows_f32; } return nullptr; + case GGML_OP_CROSS_ENTROPY_LOSS: + if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return src0->ne[0] > 1024 ? ctx->device->pipeline_cross_entropy_loss_f32_wg512 : ctx->device->pipeline_cross_entropy_loss_f32; + } + return nullptr; + case GGML_OP_CROSS_ENTROPY_LOSS_BACK: + // src0 is the scalar grad; src1 is logits + if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && src2 && src2->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + return src1->ne[0] > 1024 ? ctx->device->pipeline_cross_entropy_loss_back_f32_wg512 : ctx->device->pipeline_cross_entropy_loss_back_f32; + } + return nullptr; case GGML_OP_CUMSUM: if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { if (src0->ne[0] <= 512) { @@ -11674,6 +11779,12 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const return ctx->device->pipeline_gated_linear_attn_f32; } return nullptr; + case GGML_OP_LIGHTNING_INDEXER: + // only the k type selects a pipeline, the other types are fixed by ggml_lightning_indexer() + if (ggml_vk_lightning_indexer_k_type_supported(src1->type)) { + return ctx->device->pipeline_lightning_indexer_f32[src1->type]; + } + return nullptr; case GGML_OP_GATED_DELTA_NET: if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { const uint32_t S_v = dst->src[2]->ne[0]; @@ -12749,6 +12860,55 @@ static void ggml_vk_gated_linear_attn(ggml_backend_vk_context * ctx, vk_context& pc, { (uint32_t)(n_seqs * n_heads), 1, 1 }); } +static void ggml_vk_lightning_indexer(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * w = dst->src[2]; + const ggml_tensor * m = dst->src[3]; + + vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, q, k, w, dst, dst->op); + GGML_ASSERT(pipeline != nullptr); + + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + const uint32_t n_kv = k->ne[2]; + const uint32_t n_heads = q->ne[1]; + const uint32_t n_tokens = q->ne[2]; + const uint32_t n_streams = q->ne[3]; + const uint32_t n_masks = m->ne[3]; + + const uint32_t n_outputs = (uint32_t)(dst->ne[0] * dst->ne[1] * dst->ne[3]); + const uint32_t dispatch_x = std::min(n_outputs, ctx->device->properties.limits.maxComputeWorkGroupCount[0]); + const uint32_t dispatch_y = CEIL_DIV(n_outputs, dispatch_x); + + // q, w and dst are f32 and m is f16, so their strides are passed in elements; + // k may be quantized, so its strides stay in bytes + const uint32_t q_nb1 = q->nb[1] / sizeof(float); + const uint32_t q_nb2 = q->nb[2] / sizeof(float); + const uint32_t q_nb3 = q->nb[3] / sizeof(float); + const uint32_t k_nb2 = k->nb[2]; + const uint32_t k_nb3 = k->nb[3]; + const uint32_t w_nb1 = w->nb[1] / sizeof(float); + const uint32_t w_nb3 = w->nb[3] / sizeof(float); + const uint32_t m_nb1 = m->nb[1] / sizeof(ggml_fp16_t); + const uint32_t m_nb3 = m->nb[3] / sizeof(ggml_fp16_t); + const uint32_t d_nb1 = dst->nb[1] / sizeof(float); + const uint32_t d_nb3 = dst->nb[3] / sizeof(float); + + const vk_op_lightning_indexer_push_constants pc = { + n_kv, n_heads, n_tokens, n_streams, n_masks, dispatch_x, + q_nb1, q_nb2, q_nb3, + k_nb2, k_nb3, + w_nb1, w_nb3, + m_nb1, m_nb3, + d_nb1, d_nb3, + }; + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + {ggml_vk_tensor_subbuffer(ctx, q), ggml_vk_tensor_subbuffer(ctx, k), ggml_vk_tensor_subbuffer(ctx, w), ggml_vk_tensor_subbuffer(ctx, m), ggml_vk_tensor_subbuffer(ctx, dst)}, + pc, {dispatch_x, dispatch_y, 1}); +} + static void ggml_vk_gated_delta_net(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { const ggml_tensor * src_q = dst->src[0]; const ggml_tensor * src_v = dst->src[2]; @@ -13942,6 +14102,103 @@ static void ggml_vk_cumsum(ggml_backend_vk_context * ctx, vk_context& subctx, co ctx->prealloc_split_k_need_sync = true; } +static std::array ggml_vk_nrows_elements(uint32_t nr) { + if (nr > 262144) { + return { 512, 512, CEIL_DIV(nr, 262144) }; + } + if (nr > 512) { + return { 512, CEIL_DIV(nr, 512), 1 }; + } + return { nr, 1, 1 }; +} + +static void ggml_vk_cross_entropy_loss(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(src0)); + GGML_ASSERT(ggml_is_contiguous(src1)); + GGML_ASSERT(ggml_is_contiguous(dst)); + GGML_ASSERT(ggml_are_same_shape(src0, src1)); + GGML_ASSERT(ggml_is_scalar(dst)); + + const uint32_t nclasses = (uint32_t)src0->ne[0]; + const uint32_t nrows = (uint32_t)ggml_nrows(src0); + + vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, src0, src1, nullptr, dst, GGML_OP_CROSS_ENTROPY_LOSS); + GGML_ASSERT(pipeline != nullptr); + + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_sum_rows_f32, 1); + + vk_subbuffer src0_buf = ggml_vk_tensor_subbuffer(ctx, src0); + vk_subbuffer src1_buf = ggml_vk_tensor_subbuffer(ctx, src1); + vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst, true); + + const vk_op_push_constants pc = { nclasses, nrows, 0.0f, 0.0f, 0.0f, 0.0f }; + + const size_t tmp_size = (size_t)nrows * sizeof(float); + if (ctx->prealloc_size_x < tmp_size) { + ctx->prealloc_size_x = tmp_size; + ggml_vk_preallocate_buffers(ctx, subctx); + } + if (ctx->prealloc_x_need_sync) { + ggml_vk_sync_buffers(ctx, subctx); + } + + vk_subbuffer tmp_buf = { ctx->prealloc_x, 0, tmp_size }; + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { src0_buf, src1_buf, tmp_buf }, pc, ggml_vk_nrows_elements(nrows)); + ggml_vk_sync_buffers(ctx, subctx); + + vk_op_sum_rows_push_constants sp = {}; + sp.n_cols = nrows; + sp.ne01 = 1; + sp.ne02 = 1; + sp.weight = 1.0f; + init_pushconst_fastdiv(sp); + sp.misalign_offsets = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); + + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_sum_rows_f32, { tmp_buf, dst_buf }, sp, { 1, 1, 1 }); + ctx->prealloc_x_need_sync = true; +} + +static void ggml_vk_cross_entropy_loss_back(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { + const ggml_tensor * grad = dst->src[0]; + const ggml_tensor * logits = dst->src[1]; + const ggml_tensor * labels = dst->src[2]; + + GGML_ASSERT(grad->type == GGML_TYPE_F32); + GGML_ASSERT(logits->type == GGML_TYPE_F32); + GGML_ASSERT(labels->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_scalar(grad)); + GGML_ASSERT(ggml_is_contiguous(grad)); + GGML_ASSERT(ggml_is_contiguous(logits)); + GGML_ASSERT(ggml_is_contiguous(labels)); + GGML_ASSERT(ggml_is_contiguous(dst)); + GGML_ASSERT(ggml_are_same_shape(logits, labels)); + GGML_ASSERT(ggml_are_same_shape(logits, dst)); + + const uint32_t nclasses = (uint32_t)logits->ne[0]; + const uint32_t nrows = (uint32_t)ggml_nrows(logits); + + vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, grad, logits, labels, dst, GGML_OP_CROSS_ENTROPY_LOSS_BACK); + GGML_ASSERT(pipeline != nullptr); + + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + vk_subbuffer grad_buf = ggml_vk_tensor_subbuffer(ctx, grad); + vk_subbuffer logits_buf = ggml_vk_tensor_subbuffer(ctx, logits); + vk_subbuffer labels_buf = ggml_vk_tensor_subbuffer(ctx, labels); + vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst); + + const vk_op_push_constants pc = { nclasses, nrows, 0.0f, 0.0f, 0.0f, 0.0f }; + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { grad_buf, logits_buf, labels_buf, dst_buf }, pc, ggml_vk_nrows_elements(nrows)); +} + static void ggml_vk_argmax(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) { ggml_vk_op_f32(ctx, subctx, src0, nullptr, nullptr, nullptr, dst, GGML_OP_ARGMAX, { (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], 0.0f, 0.0f, 0.0f, 0.0f }); } @@ -15267,7 +15524,7 @@ static void ggml_vk_preallocate_buffers(ggml_backend_vk_context * ctx, vk_contex ctx->prealloc_y = ggml_vk_create_buffer_device(ctx->device, ctx->prealloc_size_y); ctx->prealloc_y_last_pipeline_used = nullptr; ctx->prealloc_y_last_tensor_used = nullptr; - ctx->prealloc_y_last_decode_vector_staging = false; + ctx->prealloc_y_last_k_padded = false; } if (ctx->prealloc_split_k == nullptr || (ctx->prealloc_size_split_k > 0 && ctx->prealloc_split_k->size < ctx->prealloc_size_split_k)) { VK_LOG_MEMORY("ggml_vk_preallocate_buffers(split_k_size: " << ctx->prealloc_size_split_k << ")"); @@ -15687,6 +15944,14 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr case GGML_OP_ARGMAX: ggml_vk_argmax(ctx, compute_ctx, src0, node); + break; + case GGML_OP_CROSS_ENTROPY_LOSS: + ggml_vk_cross_entropy_loss(ctx, compute_ctx, node); + + break; + case GGML_OP_CROSS_ENTROPY_LOSS_BACK: + ggml_vk_cross_entropy_loss_back(ctx, compute_ctx, node); + break; case GGML_OP_COUNT_EQUAL: ggml_vk_count_equal(ctx, compute_ctx, src0, src1, node); @@ -15770,6 +16035,11 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr break; + case GGML_OP_LIGHTNING_INDEXER: + ggml_vk_lightning_indexer(ctx, compute_ctx, node); + + break; + case GGML_OP_GATED_DELTA_NET: ggml_vk_gated_delta_net(ctx, compute_ctx, node); @@ -15879,7 +16149,7 @@ static void ggml_vk_graph_cleanup(ggml_backend_vk_context * ctx) { VK_LOG_DEBUG("ggml_vk_graph_cleanup()"); ctx->prealloc_y_last_pipeline_used = {}; ctx->prealloc_y_last_tensor_used = nullptr; - ctx->prealloc_y_last_decode_vector_staging = false; + ctx->prealloc_y_last_k_padded = false; ctx->unsynced_nodes_written.clear(); ctx->unsynced_nodes_read.clear(); @@ -15931,7 +16201,7 @@ static void ggml_vk_cleanup(ggml_backend_vk_context * ctx) { ctx->prealloc_y_last_pipeline_used = nullptr; ctx->prealloc_y_last_tensor_used = nullptr; - ctx->prealloc_y_last_decode_vector_staging = false; + ctx->prealloc_y_last_k_padded = false; ctx->prealloc_size_x = 0; ctx->prealloc_size_y = 0; @@ -17129,7 +17399,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->prealloc_y_last_pipeline_used = nullptr; ctx->prealloc_y_last_tensor_used = nullptr; - ctx->prealloc_y_last_decode_vector_staging = false; + ctx->prealloc_y_last_k_padded = false; if (ctx->prealloc_size_add_rms_partials) { ggml_vk_preallocate_buffers(ctx, nullptr); @@ -17534,20 +17804,32 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * return; } - auto const &is_empty = [](ggml_tensor * node) -> bool { + auto const &is_empty = [](const ggml_tensor * node) -> bool { return node->op == GGML_OP_NONE || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE; }; - auto const &is_src_of = [](const ggml_tensor *dst, const ggml_tensor *src) -> bool { + auto const &is_src_of = [&is_empty](const ggml_tensor *dst, const ggml_tensor *src) -> bool { + auto const &base = [](const ggml_tensor * tensor) { + return tensor->view_src ? tensor->view_src : tensor; + }; for (uint32_t s = 0; s < GGML_MAX_SRC; ++s) { if (dst->src[s] == src) { return true; } + if (is_empty(dst) || is_empty(src)) { + continue; + } + // A source view of dst may read storage written through a different view by src. + if (dst->src[s] && base(dst->src[s]) == base(src)) { + return true; + } + // Moving dst forward may overwrite storage still read through a view by src. + if (src->src[s] && base(dst) == base(src->src[s])) { + return true; + } } // implicit dependency if they view the same tensor - const ggml_tensor *dst2 = dst->view_src ? dst->view_src : dst; - const ggml_tensor *src2 = src->view_src ? src->view_src : src; - if (dst2 == src2) { + if (base(dst) == base(src)) { return true; } return false; @@ -18511,6 +18793,18 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm } case GGML_OP_ARGMAX: return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32; + case GGML_OP_CROSS_ENTROPY_LOSS: + return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32 + && ggml_is_contiguous(op->src[1]) && op->src[1]->type == GGML_TYPE_F32 + && ggml_are_same_shape(op->src[0], op->src[1]) + && ggml_is_contiguous(op) && ggml_is_scalar(op) && op->type == GGML_TYPE_F32; + case GGML_OP_CROSS_ENTROPY_LOSS_BACK: + return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32 && ggml_is_scalar(op->src[0]) + && ggml_is_contiguous(op->src[1]) && op->src[1]->type == GGML_TYPE_F32 + && ggml_is_contiguous(op->src[2]) && op->src[2]->type == GGML_TYPE_F32 + && ggml_are_same_shape(op->src[1], op->src[2]) + && ggml_are_same_shape(op->src[1], op) + && ggml_is_contiguous(op) && op->type == GGML_TYPE_F32; case GGML_OP_COUNT_EQUAL: return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_I32 && ggml_is_contiguous(op->src[1]) && op->src[1]->type == GGML_TYPE_I32; @@ -18536,6 +18830,40 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm case GGML_OP_GATED_LINEAR_ATTN: // the shader block size is hardcoded to head_size 64 return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && op->src[0]->ne[0] == 64; + case GGML_OP_LIGHTNING_INDEXER: + { + const ggml_tensor * q = op->src[0]; + const ggml_tensor * k = op->src[1]; + const ggml_tensor * w = op->src[2]; + const ggml_tensor * m = op->src[3]; + + // the q/w/m types and the shape relationships between q, k, w, m and dst + // are already asserted in ggml_lightning_indexer() + if (!ggml_vk_lightning_indexer_k_type_supported(k->type) || !device->fp16) { + return false; + } + + // the shader block size is hardcoded to head size 128 + if (q->ne[0] != 128) { + return false; + } + + // the shader indexes the buffers by element stride, and is dispatched + // without allow_misalign + for (const ggml_tensor * t : {q, k, w, m, op}) { + if (t->nb[0] != ggml_type_size(t->type) || + (vk_tensor_offset(t) + t->view_offs) % device->properties.limits.minStorageBufferOffsetAlignment != 0) { + return false; + } + // the strides get scaled down from bytes, so the division must be exact + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + if (t->nb[i] % ggml_type_size(t->type) != 0) { + return false; + } + } + } + return true; + } case GGML_OP_GATED_DELTA_NET: { const uint32_t S_v = op->src[2]->ne[0]; @@ -19437,6 +19765,10 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * tensor_clone = ggml_mean(ggml_ctx, src_clone[0]); } else if (tensor->op == GGML_OP_ARGMAX) { tensor_clone = ggml_argmax(ggml_ctx, src_clone[0]); + } else if (tensor->op == GGML_OP_CROSS_ENTROPY_LOSS) { + tensor_clone = ggml_cross_entropy_loss(ggml_ctx, src_clone[0], src_clone[1]); + } else if (tensor->op == GGML_OP_CROSS_ENTROPY_LOSS_BACK) { + tensor_clone = ggml_cross_entropy_loss_back(ggml_ctx, src_clone[0], src_clone[1], src_clone[2]); } else if (tensor->op == GGML_OP_COUNT_EQUAL) { tensor_clone = ggml_count_equal(ggml_ctx, src_clone[0], src_clone[1]); } else if (tensor->op == GGML_OP_SOLVE_TRI) { @@ -19541,6 +19873,8 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * const float * op_params = (const float *)tensor->op_params; tensor_clone = ggml_gated_linear_attn(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3], src_clone[4], op_params[0]); + } else if (tensor->op == GGML_OP_LIGHTNING_INDEXER) { + tensor_clone = ggml_lightning_indexer(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3]); } else if (tensor->op == GGML_OP_GATED_DELTA_NET) { tensor_clone = ggml_gated_delta_net(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3], src_clone[4], src_clone[5], diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/conv2d_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/conv2d_mm.comp index 99400098bf2..c64004cdc48 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/conv2d_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/conv2d_mm.comp @@ -19,6 +19,7 @@ #endif #include "types.glsl" +#include "utils.glsl" // shape notation: [dim(N), ..., dim(0)] -- stride(dim(j)) >= stride(dim(i)) if i > j layout(binding = 0) readonly buffer A { @@ -193,14 +194,6 @@ uint32_t Br = tid / BS_NPQ; uint32_t Bc = tid % BS_NPQ; const uint32_t BrpWg = WG_SIZE / BS_NPQ; -// see init_fastdiv_values in ggml-vulkan.cpp -uint fastdiv(uint n, uint mp, uint L) { - uint msbs, lsbs; - // msbs = mulhi(n, mp) - umulExtended(n, mp, msbs, lsbs); - return (msbs + n) >> L; -} - #ifdef COOPMAT2 #define ACC_TYPE float16_t diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/conv3d_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/conv3d_mm.comp index f66f299f6da..d5ce4290b93 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/conv3d_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/conv3d_mm.comp @@ -15,6 +15,7 @@ #endif #include "types.glsl" +#include "utils.glsl" // shape notation: [dim(N), ..., dim(0)] -- stride(dim(j)) >= stride(dim(i)) if i > j layout(binding = 0) readonly buffer A { @@ -178,14 +179,6 @@ uint32_t Br = tid / BS_NPQ; uint32_t Bc = tid % BS_NPQ; const uint32_t BrpWg = WG_SIZE / BS_NPQ; -// see init_fastdiv_values in ggml-vulkan.cpp -uint fastdiv(uint n, uint mp, uint L) { - uint msbs, lsbs; - // msbs = mulhi(n, mp) - umulExtended(n, mp, msbs, lsbs); - return (msbs + n) >> L; -} - void split_crs(uint32_t crs_idx, out uint32_t ic, out uint32_t kd, out uint32_t kh, out uint32_t kw) { const uint32_t KHKW = KH * KW; const uint32_t KDKHKW = KD * KHKW; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp b/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp index ffc8608691f..ef659959d95 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/count_experts.comp @@ -2,7 +2,13 @@ #extension GL_EXT_control_flow_attributes : enable +#ifdef USE_SUBGROUPS +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_KHR_shader_subgroup_arithmetic : enable +#endif + #include "types.glsl" +#include "utils.glsl" layout (push_constant) uniform parameter { @@ -11,6 +17,10 @@ layout (push_constant) uniform parameter uint32_t nb00; uint32_t nb01; uint32_t a_offset; + uint32_t n_experts; + uint32_t hoist_row_ids; + uint32_t ne00mp; + uint32_t ne00L; } p; #define BLOCK_SIZE 256 @@ -21,16 +31,90 @@ layout (binding = 0) readonly buffer A {uint data_a[];}; layout (binding = 1) writeonly buffer D {uint data_d[];}; shared uint vals[BLOCK_SIZE]; +shared uint offsets[BLOCK_SIZE]; +shared uint cursors[BLOCK_SIZE]; +// data_d layout when p.hoist_row_ids is set: +// [0, n_experts) per-expert row count +// [n_experts, 2*n_experts) per-expert start offset into the row id region +// [2*n_experts] total row count +// [2*n_experts + 1, ) row ids grouped by expert, packed as (i01 << 16) | (i00 & 0xffff) +// Otherwise only data_d[expert_id] is written, holding that expert's row count. void main() { const uint expert_id = gl_WorkGroupID.x; const uint num_elements = p.ne00 * p.ne01; const uint tid = gl_LocalInvocationID.x; + if (p.hoist_row_ids != 0) { + if (tid < p.n_experts) { + vals[tid] = 0; + } + barrier(); + + for (uint idx = tid; idx < num_elements; idx += BLOCK_SIZE) { + const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L); + const uint i00 = idx - i01 * p.ne00; + const uint expert = data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00]; + if (expert < p.n_experts) { + atomicAdd(vals[expert], 1); + } + } + barrier(); + +#ifdef USE_SUBGROUPS + if (gl_SubgroupID == 0) { + // pad the trip count so the subgroup ops stay in uniform control flow + const uint n_experts_padded = (p.n_experts + gl_SubgroupSize - 1) & ~(gl_SubgroupSize - 1); + uint base = 0; + for (uint expert = gl_SubgroupInvocationID; expert < n_experts_padded; expert += gl_SubgroupSize) { + const bool in_range = expert < p.n_experts; + const uint count = in_range ? vals[expert] : 0; + const uint offset = base + subgroupExclusiveAdd(count); + if (in_range) { + data_d[expert] = count; + data_d[p.n_experts + expert] = offset; + offsets[expert] = offset; + cursors[expert] = 0; + } + base += subgroupAdd(count); + } + if (subgroupElect()) { + data_d[2 * p.n_experts] = base; + } + } +#else + if (tid == 0) { + uint offset = 0; + for (uint expert = 0; expert < p.n_experts; ++expert) { + const uint count = vals[expert]; + data_d[expert] = count; + data_d[p.n_experts + expert] = offset; + offsets[expert] = offset; + cursors[expert] = 0; + offset += count; + } + data_d[2 * p.n_experts] = offset; + } +#endif + barrier(); + + for (uint idx = tid; idx < num_elements; idx += BLOCK_SIZE) { + const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L); + const uint i00 = idx - i01 * p.ne00; + const uint expert = data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00]; + if (expert < p.n_experts) { + const uint row = atomicAdd(cursors[expert], 1); + const uint packed_row_id = (i01 << 16) | (i00 & 0xffffu); + data_d[2 * p.n_experts + 1 + offsets[expert] + row] = packed_row_id; + } + } + return; + } + uint count = 0; for (uint idx = tid; idx < num_elements; idx += BLOCK_SIZE) { - const uint i01 = idx / p.ne00; - const uint i00 = idx % p.ne00; + const uint i01 = fastdiv(idx, p.ne00mp, p.ne00L); + const uint i00 = idx - i01 * p.ne00; const uint a = data_a[p.a_offset + i01 * p.nb01 + i00 * p.nb00]; count += uint(a == expert_id); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/cross_entropy_loss.comp b/ggml/src/ggml-vulkan/vulkan-shaders/cross_entropy_loss.comp new file mode 100644 index 00000000000..0c135c6fd24 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/cross_entropy_loss.comp @@ -0,0 +1,78 @@ +#version 450 + +#include "generic_head.glsl" +#include "types.glsl" + +#extension GL_EXT_control_flow_attributes : enable + +layout(constant_id = 0) const uint BLOCK_SIZE = 32; +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; +layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; +layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; + +shared FLOAT_TYPE tmp[BLOCK_SIZE]; + +FLOAT_TYPE wg_reduce_max(FLOAT_TYPE v) { + const uint tid = gl_LocalInvocationID.x; + tmp[tid] = v; + barrier(); + [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + tmp[tid] = max(tmp[tid], tmp[tid + s]); + } + barrier(); + } + v = tmp[0]; + barrier(); + return v; +} + +FLOAT_TYPE wg_reduce_sum(FLOAT_TYPE v) { + const uint tid = gl_LocalInvocationID.x; + tmp[tid] = v; + barrier(); + [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + tmp[tid] += tmp[tid + s]; + } + barrier(); + } + v = tmp[0]; + barrier(); + return v; +} + +void main() { + const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x; + const uint tid = gl_LocalInvocationID.x; + + if (row >= p.KY) { + return; + } + + const uint off = row * p.KX; + + FLOAT_TYPE max_logit = FLOAT_TYPE(uintBitsToFloat(0xFF800000)); + for (uint i = tid; i < p.KX; i += BLOCK_SIZE) { + max_logit = max(max_logit, FLOAT_TYPE(data_a[off + i])); + } + max_logit = wg_reduce_max(max_logit); + + FLOAT_TYPE sum_exp = FLOAT_TYPE(0.0f); + for (uint i = tid; i < p.KX; i += BLOCK_SIZE) { + sum_exp += exp(FLOAT_TYPE(data_a[off + i]) - max_logit); + } + const FLOAT_TYPE log_sum = log(wg_reduce_sum(sum_exp)); + + FLOAT_TYPE loss = FLOAT_TYPE(0.0f); + for (uint i = tid; i < p.KX; i += BLOCK_SIZE) { + loss += (FLOAT_TYPE(data_a[off + i]) - max_logit - log_sum) * FLOAT_TYPE(data_b[off + i]); + } + loss = -wg_reduce_sum(loss) / FLOAT_TYPE(p.KY); + + if (tid == 0) { + data_d[row] = D_TYPE(loss); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/cross_entropy_loss_back.comp b/ggml/src/ggml-vulkan/vulkan-shaders/cross_entropy_loss_back.comp new file mode 100644 index 00000000000..3cdebe86e47 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/cross_entropy_loss_back.comp @@ -0,0 +1,75 @@ +#version 450 + +#include "generic_head.glsl" +#include "types.glsl" + +#extension GL_EXT_control_flow_attributes : enable + +layout(constant_id = 0) const uint BLOCK_SIZE = 32; +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer G {A_TYPE data_g[];}; +layout (binding = 1) readonly buffer X {B_TYPE data_x[];}; +layout (binding = 2) readonly buffer Y {B_TYPE data_y[];}; +layout (binding = 3) writeonly buffer D {D_TYPE data_d[];}; + +shared FLOAT_TYPE tmp[BLOCK_SIZE]; + +FLOAT_TYPE wg_reduce_max(FLOAT_TYPE v) { + const uint tid = gl_LocalInvocationID.x; + tmp[tid] = v; + barrier(); + [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + tmp[tid] = max(tmp[tid], tmp[tid + s]); + } + barrier(); + } + v = tmp[0]; + barrier(); + return v; +} + +FLOAT_TYPE wg_reduce_sum(FLOAT_TYPE v) { + const uint tid = gl_LocalInvocationID.x; + tmp[tid] = v; + barrier(); + [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + tmp[tid] += tmp[tid + s]; + } + barrier(); + } + v = tmp[0]; + barrier(); + return v; +} + +void main() { + const uint row = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x; + const uint tid = gl_LocalInvocationID.x; + + if (row >= p.KY) { + return; + } + + const uint off = row * p.KX; + const FLOAT_TYPE d_by_nrows = FLOAT_TYPE(data_g[0]) / FLOAT_TYPE(p.KY); + + FLOAT_TYPE max_logit = FLOAT_TYPE(uintBitsToFloat(0xFF800000)); + for (uint i = tid; i < p.KX; i += BLOCK_SIZE) { + max_logit = max(max_logit, FLOAT_TYPE(data_x[off + i])); + } + max_logit = wg_reduce_max(max_logit); + + FLOAT_TYPE sum_exp = FLOAT_TYPE(0.0f); + for (uint i = tid; i < p.KX; i += BLOCK_SIZE) { + sum_exp += exp(FLOAT_TYPE(data_x[off + i]) - max_logit); + } + const FLOAT_TYPE inv_sum = FLOAT_TYPE(1.0f) / wg_reduce_sum(sum_exp); + + for (uint i = tid; i < p.KX; i += BLOCK_SIZE) { + const FLOAT_TYPE sm = exp(FLOAT_TYPE(data_x[off + i]) - max_logit) * inv_sum; + data_d[off + i] = D_TYPE((sm - FLOAT_TYPE(data_y[off + i])) * d_by_nrows); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl new file mode 100644 index 00000000000..6f414ded123 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl @@ -0,0 +1,55 @@ +#if !defined(GGML_FA_TYPES_COMP) +#define GGML_FA_TYPES_COMP + +// FaTypeK / FaTypeV spec constant values. These mirror enum ggml_type so the +// host can pass the type directly. Keep in sync with ggml.h. +#define FA_TYPE_F32 0u +#define FA_TYPE_F16 1u +#define FA_TYPE_Q4_0 2u +#define FA_TYPE_Q4_1 3u +#define FA_TYPE_Q5_0 6u +#define FA_TYPE_Q5_1 7u +#define FA_TYPE_Q8_0 8u +#define FA_TYPE_IQ4_NL 20u +#define FA_TYPE_BF16 30u + +// Number of matrix elements per buffer block, derived from the K/V type spec +// constant. F32 is treated as a vec4 "block" of 4 floats. F16 uses block size 1 +// and bypasses the dequant path entirely. Quants follow their ggml block sizes. +uint fa_block_elems(uint ty) { + switch (ty) { + case FA_TYPE_F32: return 4u; + case FA_TYPE_F16: return 1u; + case FA_TYPE_Q4_0: return uint(QUANT_K_Q4_0); + case FA_TYPE_Q4_1: return uint(QUANT_K_Q4_1); + case FA_TYPE_Q5_0: return uint(QUANT_K_Q5_0); + case FA_TYPE_Q5_1: return uint(QUANT_K_Q5_1); + case FA_TYPE_Q8_0: return uint(QUANT_K_Q8_0); + case FA_TYPE_IQ4_NL: return uint(QUANT_K_IQ4_NL); + case FA_TYPE_BF16: return 1u; + default: return 1u; + } +} + +// QUANT_R_MMQ for FA-eligible K types. Q4_*/Q5_* store two nibbles per byte +// (R==2); Q8_0 stores one byte per element (R==1). Used to derive the number +// of int32s per 32-element block on the MMQ K path: ints_per_block == 8 / R. +uint fa_quant_r_mmq(uint ty) { + switch (ty) { + case FA_TYPE_Q4_0: return uint(QUANT_R_Q4_0); + case FA_TYPE_Q4_1: return uint(QUANT_R_Q4_1); + case FA_TYPE_Q5_0: return uint(QUANT_R_Q5_0); + case FA_TYPE_Q5_1: return uint(QUANT_R_Q5_1); + case FA_TYPE_Q8_0: return uint(QUANT_R_Q8_0); + default: return 1u; + } +} + +bool fa_type_needs_shmem(uint ty) { + switch (ty) { + case FA_TYPE_IQ4_NL: return true; + default: return false; + } +} + +#endif // !defined(GGML_FA_TYPES_COMP) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl index 3c64f91dad3..0ce4503a884 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl @@ -88,17 +88,7 @@ layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];}; #define BINDING_IDX_K 0 #define BINDING_IDX_V 1 -// FaTypeK / FaTypeV spec constant values. These mirror enum ggml_type so the -// host can pass the type directly. Keep in sync with ggml.h. -#define FA_TYPE_F32 0u -#define FA_TYPE_F16 1u -#define FA_TYPE_Q4_0 2u -#define FA_TYPE_Q4_1 3u -#define FA_TYPE_Q5_0 6u -#define FA_TYPE_Q5_1 7u -#define FA_TYPE_Q8_0 8u -#define FA_TYPE_IQ4_NL 20u -#define FA_TYPE_BF16 30u +#include "fa_types.glsl" #if defined(BFLOAT16) #define O_TYPE float @@ -108,45 +98,6 @@ layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];}; #define O_TYPEV4 FLOAT_TYPEV4 #endif -// Number of matrix elements per buffer block, derived from the K/V type spec -// constant. F32 is treated as a vec4 "block" of 4 floats. F16 uses block size 1 -// and bypasses the dequant path entirely. Quants follow their ggml block sizes. -uint fa_block_elems(uint ty) { - switch (ty) { - case FA_TYPE_F32: return 4u; - case FA_TYPE_F16: return 1u; - case FA_TYPE_Q4_0: return uint(QUANT_K_Q4_0); - case FA_TYPE_Q4_1: return uint(QUANT_K_Q4_1); - case FA_TYPE_Q5_0: return uint(QUANT_K_Q5_0); - case FA_TYPE_Q5_1: return uint(QUANT_K_Q5_1); - case FA_TYPE_Q8_0: return uint(QUANT_K_Q8_0); - case FA_TYPE_IQ4_NL: return uint(QUANT_K_IQ4_NL); - case FA_TYPE_BF16: return 1u; - default: return 1u; - } -} - -// QUANT_R_MMQ for FA-eligible K types. Q4_*/Q5_* store two nibbles per byte -// (R==2); Q8_0 stores one byte per element (R==1). Used to derive the number -// of int32s per 32-element block on the MMQ K path: ints_per_block == 8 / R. -uint fa_quant_r_mmq(uint ty) { - switch (ty) { - case FA_TYPE_Q4_0: return uint(QUANT_R_Q4_0); - case FA_TYPE_Q4_1: return uint(QUANT_R_Q4_1); - case FA_TYPE_Q5_0: return uint(QUANT_R_Q5_0); - case FA_TYPE_Q5_1: return uint(QUANT_R_Q5_1); - case FA_TYPE_Q8_0: return uint(QUANT_R_Q8_0); - default: return 1u; - } -} - -bool fa_type_needs_shmem(uint ty) { - switch (ty) { - case FA_TYPE_IQ4_NL: return true; - default: return false; - } -} - // These can't be `const` globals because GLSL forbids function calls in global // const initializers, even when the spec constants would let the driver fold // them. Macros expand at the use site and fold after specialization. diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.glsl index 9d4176f3f96..e13de9a00f2 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/generic_unary_head.glsl @@ -1,6 +1,8 @@ #extension GL_EXT_shader_16bit_storage : require #extension GL_EXT_control_flow_attributes : require +#include "utils.glsl" + layout (push_constant) uniform parameter { uint ne; @@ -32,18 +34,6 @@ uint get_idx() { uint get_aoffset() { return p.misalign_offsets >> 16; } uint get_doffset() { return p.misalign_offsets & 0xFFFF; } -// see init_fastdiv_values in ggml-vulkan.cpp -uint fastdiv(uint n, uint mp, uint L) { - uint msbs, lsbs; - // msbs = mulhi(n, mp) - umulExtended(n, mp, msbs, lsbs); - return (msbs + n) >> L; -} - -uint fastdiv_L(uint packed, uint slot) { - return (packed >> (slot * 8)) & 0x3Fu; -} - uint src0_idx(uint idx) { const uint i03 = fastdiv(idx, p.ne0_012mp, fastdiv_L(p.ne0_Ls, 0)); const uint i03_offset = i03 * p.ne02*p.ne01*p.ne00; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/glu_head.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/glu_head.glsl index c3cae736f97..fc2951ec2e5 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/glu_head.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/glu_head.glsl @@ -1,5 +1,7 @@ #extension GL_EXT_shader_16bit_storage : require +#include "utils.glsl" + layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; @@ -39,9 +41,3 @@ uint get_aoffset() { return p.misalign_offsets >> 16; } uint get_boffset() { return (p.misalign_offsets >> 8) & 0xFF; } uint get_doffset() { return p.misalign_offsets & 0xFF; } -// see init_fastdiv_values in ggml-vulkan.cpp -uint fastdiv(uint n, uint mp, uint L) { - uint msbs, lsbs; - umulExtended(n, mp, msbs, lsbs); - return (msbs + n) >> L; -} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp new file mode 100644 index 00000000000..ba76ec72ca6 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp @@ -0,0 +1,151 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_KHR_shader_subgroup_basic : enable +#if USE_SUBGROUP_ADD +#extension GL_KHR_shader_subgroup_arithmetic : enable +#endif + +#define BINDING_IDX_K 0u + +#include "types.glsl" +#include "fa_types.glsl" +#define FaTypeV FA_TYPE_F32 + +layout(constant_id = 0) const uint FaTypeK = FA_TYPE_F32; +layout(constant_id = 1) const uint FaBlockBytesK = 4; +layout(constant_id = 2) const uint SUBGROUP_SIZE = 32; + +#include "flash_attn_dequant.glsl" + +// one workgroup computes one output element, one invocation per head element +#define HEAD_SIZE 128 + +layout(local_size_x = HEAD_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer QBuf { float q[]; }; +layout(binding = 1) readonly buffer KBufF16 { float16_t k_f16[]; }; +layout(binding = 1) readonly buffer KBufF32 { float k_f32[]; }; +layout(binding = 1) readonly buffer KBufBF16 { uint16_t k_bf16[]; }; +layout(binding = 2) readonly buffer WBuf { float weights[]; }; +layout(binding = 3) readonly buffer MBuf { float16_t mask[]; }; +layout(binding = 4) writeonly buffer DstBuf { float dst[]; }; + +layout(push_constant) uniform PushConstants { + uint n_kv; + uint n_heads; + uint n_tokens; + uint n_streams; + uint n_masks; + uint dispatch_x; + uint q_nb1; + uint q_nb2; + uint q_nb3; + uint k_nb2; + uint k_nb3; + uint w_nb1; + uint w_nb3; + uint m_nb1; + uint m_nb3; + uint d_nb1; + uint d_nb3; +}; + +shared float k_row[HEAD_SIZE]; + +#if USE_SUBGROUP_ADD +shared float sg_partials[HEAD_SIZE / SUBGROUP_SIZE]; +#else +shared float partials[HEAD_SIZE]; +#endif + +void main() { + const uint tid = gl_LocalInvocationID.x; + const uint output_idx = gl_WorkGroupID.y * dispatch_x + gl_WorkGroupID.x; + const uint n_outputs = n_kv * n_tokens * n_streams; + + if (fa_type_needs_shmem(FaTypeK)) { + init_iq_shmem(gl_WorkGroupSize); + } + + if (output_idx >= n_outputs) { + return; + } + + const uint ik = output_idx % n_kv; + const uint ts = output_idx / n_kv; + const uint t = ts % n_tokens; + const uint s = ts / n_tokens; + const uint k_offset = ik * k_nb2 + s * k_nb3; + + // k strides come in as bytes, so scale them down to the view being indexed + const uint k_block_elems = fa_block_elems(FaTypeK); + const uint k_elem_bytes = FaBlockBytesK / k_block_elems; + + if (FaTypeK == FA_TYPE_F16) { + k_row[tid] = float(k_f16[k_offset / k_elem_bytes + tid]); + } else if (FaTypeK == FA_TYPE_F32) { + k_row[tid] = k_f32[k_offset / k_elem_bytes + tid]; + } else if (FaTypeK == FA_TYPE_BF16) { + k_row[tid] = bf16_to_fp32(uint(k_bf16[k_offset / k_elem_bytes + tid])); + } else if (4 * tid < HEAD_SIZE) { + const uint coord = 4 * tid; + const uint ib = coord / k_block_elems; + const uint iqs = coord % k_block_elems; + const vec4 values = dequantize4(ib, iqs, k_offset / FaBlockBytesK, BINDING_IDX_K); + k_row[coord + 0] = values.x; + k_row[coord + 1] = values.y; + k_row[coord + 2] = values.z; + k_row[coord + 3] = values.w; + } + barrier(); + + const float k_val = k_row[tid]; + + float score = 0.0; + for (uint h = 0; h < n_heads; ++h) { + const float prod = q[h * q_nb1 + t * q_nb2 + s * q_nb3 + tid] * k_val; + +#if USE_SUBGROUP_ADD + const float sg_sum = subgroupAdd(prod); + if (gl_SubgroupInvocationID == 0) { + sg_partials[gl_SubgroupID] = sg_sum; + } + barrier(); + + if (tid == 0) { + float sum = 0.0; + [[unroll]] for (uint i = 0; i < HEAD_SIZE / SUBGROUP_SIZE; ++i) { + sum += sg_partials[i]; + } + score += max(sum, 0.0) * weights[h + t * w_nb1 + s * w_nb3]; + } + // the reads above must complete before the next iteration overwrites sg_partials + barrier(); +#else + partials[tid] = prod; + barrier(); + + [[unroll]] for (uint stride = HEAD_SIZE / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + partials[tid] += partials[tid + stride]; + } + barrier(); + } + + if (tid == 0) { + score += max(partials[0], 0.0) * weights[h + t * w_nb1 + s * w_nb3]; + } + // the read of partials[0] above must complete before the next iteration + // overwrites partials[tid] + barrier(); +#endif + } + + if (tid == 0) { + const uint mask_offset = ik + t * m_nb1 + (s % n_masks) * m_nb3; + dst[ik + t * d_nb1 + s * d_nb3] = score + float(mask[mask_offset]); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp index 3df88044a5e..63c4aaebcb1 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp @@ -88,6 +88,8 @@ layout (push_constant) uniform parameter uint nei1; uint nbi1; uint ne11; + uint n_experts; + uint hoist_row_ids; #else uint base_work_group_z; uint num_batches; @@ -214,27 +216,31 @@ void main() { const uint loadstride_b = gl_WorkGroupSize.x * LOAD_VEC_B_EFF * LOAD_VEC_BATCH_B / BK; #ifdef MUL_MAT_ID -#ifdef MUL_MAT_ID_USE_SUBGROUPS - if (bitCount(p.nei0) == 1) { - load_row_ids(expert_idx, true, ic); + if (p.hoist_row_ids != 0) { + load_row_ids_hoisted(expert_idx, ic); } else { - load_row_ids(expert_idx, false, ic); - } +#ifdef MUL_MAT_ID_USE_SUBGROUPS + if (bitCount(p.nei0) == 1) { + load_row_ids(expert_idx, true, ic); + } else { + load_row_ids(expert_idx, false, ic); + } #else - _ne1 = 0; - for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) { - for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) { - if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) { - if (_ne1 >= ic * BN) { - row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1); + _ne1 = 0; + for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) { + for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) { + if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) { + if (_ne1 >= ic * BN) { + row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1); + } + _ne1++; } - _ne1++; } } - } - barrier(); + barrier(); #endif + } // Workgroup has no work if (ic * BN >= _ne1) return; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp index a2e15f6f5ce..27f3178e7f2 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp @@ -56,6 +56,8 @@ layout (push_constant) uniform parameter uint nei1; uint nbi1; uint ne11; + uint n_experts; + uint hoist_row_ids; #else uint base_work_group_z; uint num_batches; @@ -64,9 +66,9 @@ layout (push_constant) uniform parameter uint ne12; uint broadcast2; uint broadcast3; -#endif // N dimension for the B matrix can be >= p.N uint padded_N; +#endif } p; @@ -225,6 +227,23 @@ void load_row_ids(uint expert_idx, bool nei0_is_pow2, uint ic) { } barrier(); } + +void load_row_ids_hoisted(uint expert_idx, uint ic) { + _ne1 = uint(data_expert_count[expert_idx]); + + const uint tile_begin = ic * BN; + const uint tile_count = tile_begin < _ne1 ? min(BN, _ne1 - tile_begin) : 0; + const uint expert_offset = uint(data_expert_count[p.n_experts + expert_idx]); + const uint row_ids_offset = 2 * p.n_experts + 1 + expert_offset + tile_begin; + + for (uint i = gl_LocalInvocationIndex; i < tile_count; i += BLOCK_SIZE) { + const uint packed_row_id = uint(data_expert_count[row_ids_offset + i]); + const uint ii0 = packed_row_id & 0xffffu; + const uint ii1 = packed_row_id >> 16; + row_ids[i] = u16vec4(fastmod(ii0, p.ne11), ii1, ii0, 0); + } + barrier(); +} #endif void main() { @@ -266,7 +285,9 @@ void main() { const uint ik = gl_WorkGroupID.x / blocks_m; #ifdef MUL_MAT_ID - if (bitCount(p.nei0) == 1) { + if (p.hoist_row_ids != 0) { + load_row_ids_hoisted(expert_idx, ic); + } else if (bitCount(p.nei0) == 1) { load_row_ids(expert_idx, true, ic); } else { load_row_ids(expert_idx, false, ic); @@ -309,7 +330,9 @@ void main() { tensorLayoutNV<2> tensorLayoutA = createTensorLayoutNV(2); tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutAClamp = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); tensorLayoutNV<2> tensorLayoutB = createTensorLayoutNV(2); +#ifndef MUL_MAT_ID tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutBClamp = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); +#endif tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutD = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); #if QUANT_K > 1 @@ -322,12 +345,19 @@ void main() { // Use end_k rather than p.K as the dimension because that's what // we need to bound check against when using split_k. - // Bounds check B against padded_N, but bounds check D against N. tensorLayoutA = setTensorLayoutDimensionNV(tensorLayoutA, p.M, end_k); +#ifdef MUL_MAT_ID + // MUL_MAT_ID pads each B row to stride_b so partial K tiles read zeros without clamping. + tensorLayoutB = setTensorLayoutDimensionNV(tensorLayoutB, BN, p.stride_b); +#else + // Bounds check B against padded_N, but bounds check D against N. tensorLayoutB = setTensorLayoutDimensionNV(tensorLayoutB, p.padded_N, end_k); +#endif tensorLayoutD = setTensorLayoutDimensionNV(tensorLayoutD, p.N, p.M); tensorLayoutAClamp = setTensorLayoutDimensionNV(tensorLayoutAClamp, p.M, end_k); +#ifndef MUL_MAT_ID tensorLayoutBClamp = setTensorLayoutDimensionNV(tensorLayoutBClamp, p.padded_N, end_k); +#endif tensorLayoutD = setTensorLayoutStrideNV(tensorLayoutD, p.stride_d, 1); @@ -504,7 +534,9 @@ void main() { tensorLayoutB = setTensorLayoutStrideNV(tensorLayoutB, stride_b, 1); +#ifndef MUL_MAT_ID tensorLayoutBClamp = setTensorLayoutStrideNV(tensorLayoutBClamp, stride_b, 1); +#endif uint k_iters = (end_k - start_k + BK - 1) / BK; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_id_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_id_funcs.glsl index 26c5c12a49a..54ad60b2efb 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_id_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_id_funcs.glsl @@ -71,4 +71,19 @@ void load_row_ids(uint expert_idx, bool nei0_is_pow2, uint ic) { barrier(); } #endif // MUL_MAT_ID_USE_SUBGROUPS + +void load_row_ids_hoisted(uint expert_idx, uint ic) { + _ne1 = uint(data_expert_count[expert_idx]); + + const uint tile_begin = ic * BN; + const uint tile_count = tile_begin < _ne1 ? min(BN, _ne1 - tile_begin) : 0; + const uint expert_offset = uint(data_expert_count[p.n_experts + expert_idx]); + const uint row_ids_offset = 2 * p.n_experts + 1 + expert_offset + tile_begin; + + for (uint i = gl_LocalInvocationIndex; i < tile_count; i += BLOCK_SIZE) { + const uint packed_row_id = uint(data_expert_count[row_ids_offset + i]); + row_ids[i] = u16vec2(packed_row_id & 0xffffu, packed_row_id >> 16); + } + barrier(); +} #endif // MUL_MAT_ID diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp index aae1c2e8ae9..1fbcbf6c933 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp @@ -56,6 +56,8 @@ layout (push_constant) uniform parameter uint nei1; uint nbi1; uint ne11; + uint n_experts; + uint hoist_row_ids; #else uint base_work_group_z; uint num_batches; @@ -157,27 +159,31 @@ void main() { const uint loadstride_b = BLOCK_SIZE * LOAD_VEC_B / BK; #ifdef MUL_MAT_ID -#ifdef MUL_MAT_ID_USE_SUBGROUPS - if (bitCount(p.nei0) == 1) { - load_row_ids(expert_idx, true, ic); + if (p.hoist_row_ids != 0) { + load_row_ids_hoisted(expert_idx, ic); } else { - load_row_ids(expert_idx, false, ic); - } +#ifdef MUL_MAT_ID_USE_SUBGROUPS + if (bitCount(p.nei0) == 1) { + load_row_ids(expert_idx, true, ic); + } else { + load_row_ids(expert_idx, false, ic); + } #else - _ne1 = 0; - for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) { - for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) { - if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) { - if (_ne1 >= ic * BN) { - row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1); + _ne1 = 0; + for (uint ii1 = 0; ii1 < p.nei1 && _ne1 < (ic + 1) * BN; ii1++) { + for (uint ii0 = 0; ii0 < p.nei0 && _ne1 < (ic + 1) * BN; ii0++) { + if (data_ids[ii1*p.nbi1 + ii0] == expert_idx) { + if (_ne1 >= ic * BN) { + row_ids[_ne1 - ic * BN] = u16vec2(ii0, ii1); + } + _ne1++; } - _ne1++; } } - } - barrier(); + barrier(); #endif + } // Workgroup has no work if (ic * BN >= _ne1) return; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.glsl index 2b841baa6bf..1cb0f7827a3 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/sum_rows.glsl @@ -1,4 +1,6 @@ +#include "utils.glsl" + // vk_op_sum_rows_push_constants layout (push_constant) uniform parameter { @@ -15,11 +17,3 @@ layout (push_constant) uniform parameter uint get_aoffset() { return p.misalign_offsets >> 16; } uint get_doffset() { return p.misalign_offsets & 0xFFFF; } -// see init_fastdiv_values in ggml-vulkan.cpp -uint fastdiv(uint n, uint mp, uint L) { - uint msbs, lsbs; - // msbs = mulhi(n, mp) - umulExtended(n, mp, msbs, lsbs); - return (msbs + n) >> L; -} - diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/utils.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/utils.glsl index dc4a1e6d96b..8aac64d7593 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/utils.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/utils.glsl @@ -9,14 +9,26 @@ uint fastmod(uint a, uint b) { return a % b; } -uint fastdiv(uint a, uint b) { +// see init_fastdiv_values in ggml-vulkan.cpp +uint fastdiv(uint n, uint mp, uint L) { + uint msbs, lsbs; + // msbs = mulhi(n, mp) + umulExtended(n, mp, msbs, lsbs); + return (msbs + n) >> L; +} + +uint fastdiv_L(uint packed, uint slot) { + return (packed >> (slot * 8)) & 0x3Fu; +} + +uint fastdiv_small(uint a, uint b) { return (a < b) ? 0 : (a / b); } void get_indices(uint idx, out uint i00, out uint i01, out uint i02, out uint i03, uint ne00, uint ne01, uint ne02, uint ne03) { - i03 = fastdiv(idx, (ne02*ne01*ne00)); + i03 = fastdiv_small(idx, (ne02*ne01*ne00)); const uint i03_offset = i03 * ne02*ne01*ne00; - i02 = fastdiv((idx - i03_offset), (ne01*ne00)); + i02 = fastdiv_small((idx - i03_offset), (ne01*ne00)); const uint i02_offset = i02*ne01*ne00; i01 = (idx - i03_offset - i02_offset) / ne00; i00 = idx - i03_offset - i02_offset - i01*ne00; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index 17d57d5a18f..d375c2d1277 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -1029,6 +1029,8 @@ void process_shaders() { string_to_spv("argmax_f32", "argmax.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "int"}})); string_to_spv("sum_rows_f32", "sum_rows.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); + string_to_spv("cross_entropy_loss_f32", "cross_entropy_loss.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}})); + string_to_spv("cross_entropy_loss_back_f32", "cross_entropy_loss_back.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("fwht_f32", "fwht.comp", {}); string_to_spv("fwht_shmem_f32", "fwht.comp", {{"FWHT_SHMEM", "1"}}); string_to_spv("count_equal_i32", "count_equal.comp", merge_maps(base_dict, {{"A_TYPE", "int"}, {"B_TYPE", "int"}, {"D_TYPE", "int"}})); @@ -1037,6 +1039,7 @@ void process_shaders() { string_to_spv("cumsum_multipass2_f32", "cumsum_multipass2.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("count_experts", "count_experts.comp", merge_maps(base_dict, {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}})); + string_to_spv("count_experts_subgroup", "count_experts.comp", merge_maps(base_dict, {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}, {"USE_SUBGROUPS", "1"}})); for (std::string dim_str : {"", "_3d"}) { for (bool bda : {false, true}) { @@ -1067,6 +1070,12 @@ void process_shaders() { string_to_spv("gated_linear_attn_f32", "gla.comp", merge_maps(base_dict, {{"A_TYPE", "float"}})); + // Compile IQ4_NL support in so its shared LUT is available when K uses it. + // K quant type is selected at runtime via the FaTypeK spec constant. + std::map li_dict = {{"FLOAT_TYPE", "float"}, {"FLOAT_TYPEV4", "vec4"}, {"DATA_A_IQ4_NL", "1"}}; + string_to_spv("lightning_indexer_f32", "lightning_indexer.comp", li_dict); + string_to_spv("lightning_indexer_subgroup_f32", "lightning_indexer.comp", merge_maps(li_dict, {{"USE_SUBGROUP_ADD", "1"}})); + string_to_spv("rwkv_wkv7_f32", "wkv7.comp", merge_maps(base_dict, {{"A_TYPE", "float"}})); string_to_spv("gated_delta_net_f32", "gated_delta_net.comp", merge_maps(base_dict, {{"FLOAT_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}, {"USE_SUBGROUP_CLUSTERED", "1"}})); diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index f236a5d2c98..c99feb3c795 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -162,7 +162,12 @@ class LLM: TARGET_LAYERS = "{arch}.target_layers" TARGET_HIDDEN_SIZE = "{arch}.target_hidden_size" BLOCK_SIZE = "{arch}.block_size" + CONV_KERNEL_SIZE = "{arch}.conv_kernel_size" + CONV_GROUP_SIZE = "{arch}.conv_group_size" + SELECTOR_RANK = "{arch}.selector_rank" + SELECTOR_TOP_K = "{arch}.selector_top_k" SAMPLE_FROM_ANCHOR = "{arch}.sample_from_anchor" + HAS_CONFIDENCE_HEAD = "{arch}.has_confidence_head" NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual" NORM_BEFORE_FC = "{arch}.norm_before_fc" @@ -225,6 +230,19 @@ class HyperConnection: COUNT = "{arch}.hyper_connection.count" SINKHORN_ITERATIONS = "{arch}.hyper_connection.sinkhorn_iterations" EPSILON = "{arch}.hyper_connection.epsilon" + # absent means the mix projection is full rank (DeepSeek-V4 behaviour) + LOW_RANK = "{arch}.hyper_connection.low_rank" + + class PerLayerEmbedding: + LAYERS = "{arch}.ple.layers" + NGRAM_SIZE = "{arch}.ple.ngram_size" + HEADS_PER_NGRAM = "{arch}.ple.heads_per_ngram" + CONV_KERNEL = "{arch}.ple.conv_kernel" + LAYER_MULTIPLIERS = "{arch}.ple.layer_multipliers" + HEAD_OFFSETS = "{arch}.ple.head_offsets" + HEAD_VOCAB_SIZES = "{arch}.ple.head_vocab_sizes" + EOS_TOKEN_ID = "{arch}.ple.eos_token_id" + IMAGE_TOKEN_ID = "{arch}.ple.image_token_id" class Rope: DIMENSION_COUNT = "{arch}.rope.dimension_count" @@ -494,6 +512,7 @@ class MODEL_ARCH(IntEnum): QWEN3VLMOE = auto() QWEN35 = auto() QWEN35MOE = auto() + QWEN4EXP = auto() PHI2 = auto() PHI3 = auto() PHIMOE = auto() @@ -636,6 +655,9 @@ class MODEL_TENSOR(IntEnum): HC_HEAD_FN = auto() HC_HEAD_BASE = auto() HC_HEAD_SCALE = auto() + HC_HEAD_NORM = auto() # qwen4exp + HC_HEAD_DOWN = auto() # qwen4exp + HC_HEAD_UP = auto() # qwen4exp ROPE_FREQS = auto() ROPE_FACTORS_LONG = auto() ROPE_FACTORS_SHORT = auto() @@ -780,6 +802,20 @@ class MODEL_TENSOR(IntEnum): HC_FFN_FN = auto() HC_FFN_BASE = auto() HC_FFN_SCALE = auto() + HC_ATTN_NORM = auto() # qwen4exp + HC_ATTN_DOWN = auto() # qwen4exp + HC_ATTN_UP = auto() # qwen4exp + HC_ATTN_INJECT = auto() # qwen4exp + HC_FFN_NORM = auto() # qwen4exp + HC_FFN_DOWN = auto() # qwen4exp + HC_FFN_UP = auto() # qwen4exp + HC_FFN_INJECT = auto() # qwen4exp + PLE_KEY = auto() # qwen4exp + PLE_VALUE = auto() # qwen4exp + PLE_NORM_KEY = auto() # qwen4exp + PLE_NORM_QUERY = auto() # qwen4exp + PLE_NORM_CONV = auto() # qwen4exp + PLE_CONV1D = auto() # qwen4exp ATTN_COMPRESSOR_WKV = auto() ATTN_COMPRESSOR_WGATE = auto() ATTN_COMPRESSOR_APE = auto() @@ -1146,6 +1182,13 @@ class MODEL_TENSOR(IntEnum): DSPARK_MARKOV_W1 = auto() # markov head: prev-token embed DSPARK_MARKOV_W2 = auto() # markov head: bias projection DSPARK_CONF_PROJ = auto() # confidence head + DFLASH_ATTN_CONV_BASE = auto() + DFLASH_ATTN_CONV_PROJ = auto() + DFLASH_FFN_CONV_BASE = auto() + DFLASH_FFN_CONV_PROJ = auto() + DFLASH_SELECTOR_PREV = auto() + DFLASH_SELECTOR_NEXT = auto() + DFLASH_SELECTOR_HIDDEN = auto() # lfm2 audio A_ENC_NORM_CONV = auto() A_ENC_LINEAR_POS = auto() @@ -1217,6 +1260,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.QWEN3VLMOE: "qwen3vlmoe", MODEL_ARCH.QWEN35: "qwen35", MODEL_ARCH.QWEN35MOE: "qwen35moe", + MODEL_ARCH.QWEN4EXP: "qwen4exp", MODEL_ARCH.PHI2: "phi2", MODEL_ARCH.PHI3: "phi3", MODEL_ARCH.PHIMOE: "phimoe", @@ -1358,6 +1402,9 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.HC_HEAD_FN: "output_hc_fn", MODEL_TENSOR.HC_HEAD_BASE: "output_hc_base", MODEL_TENSOR.HC_HEAD_SCALE: "output_hc_scale", + MODEL_TENSOR.HC_HEAD_NORM: "output_hc_norm", # qwen4exp + MODEL_TENSOR.HC_HEAD_DOWN: "output_hc_down", # qwen4exp + MODEL_TENSOR.HC_HEAD_UP: "output_hc_up", # qwen4exp MODEL_TENSOR.ROPE_FREQS: "rope_freqs", MODEL_TENSOR.ROPE_FACTORS_LONG: "rope_factors_long", MODEL_TENSOR.ROPE_FACTORS_SHORT: "rope_factors_short", @@ -1502,6 +1549,20 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.HC_FFN_FN: "blk.{bid}.hc_ffn_fn", MODEL_TENSOR.HC_FFN_BASE: "blk.{bid}.hc_ffn_base", MODEL_TENSOR.HC_FFN_SCALE: "blk.{bid}.hc_ffn_scale", + MODEL_TENSOR.HC_ATTN_NORM: "blk.{bid}.hc_attn_norm", # qwen4exp + MODEL_TENSOR.HC_ATTN_DOWN: "blk.{bid}.hc_attn_down", # qwen4exp + MODEL_TENSOR.HC_ATTN_UP: "blk.{bid}.hc_attn_up", # qwen4exp + MODEL_TENSOR.HC_ATTN_INJECT: "blk.{bid}.hc_attn_inject", # qwen4exp + MODEL_TENSOR.HC_FFN_NORM: "blk.{bid}.hc_ffn_norm", # qwen4exp + MODEL_TENSOR.HC_FFN_DOWN: "blk.{bid}.hc_ffn_down", # qwen4exp + MODEL_TENSOR.HC_FFN_UP: "blk.{bid}.hc_ffn_up", # qwen4exp + MODEL_TENSOR.HC_FFN_INJECT: "blk.{bid}.hc_ffn_inject", # qwen4exp + MODEL_TENSOR.PLE_KEY: "blk.{bid}.ple_key", # qwen4exp + MODEL_TENSOR.PLE_VALUE: "blk.{bid}.ple_value", # qwen4exp + MODEL_TENSOR.PLE_NORM_KEY: "blk.{bid}.ple_norm_key", # qwen4exp + MODEL_TENSOR.PLE_NORM_QUERY: "blk.{bid}.ple_norm_query", # qwen4exp + MODEL_TENSOR.PLE_NORM_CONV: "blk.{bid}.ple_norm_conv", # qwen4exp + MODEL_TENSOR.PLE_CONV1D: "blk.{bid}.ple_conv1d", # qwen4exp MODEL_TENSOR.ATTN_COMPRESSOR_WKV: "blk.{bid}.attn_compressor_kv", MODEL_TENSOR.ATTN_COMPRESSOR_WGATE: "blk.{bid}.attn_compressor_gate", MODEL_TENSOR.ATTN_COMPRESSOR_APE: "blk.{bid}.attn_compressor_ape", @@ -1895,6 +1956,13 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.DSPARK_MARKOV_W1: "markov_w1", MODEL_TENSOR.DSPARK_MARKOV_W2: "markov_w2", MODEL_TENSOR.DSPARK_CONF_PROJ: "conf_proj", + MODEL_TENSOR.DFLASH_ATTN_CONV_BASE: "blk.{bid}.attn_conv_base", + MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ: "blk.{bid}.attn_conv_proj", + MODEL_TENSOR.DFLASH_FFN_CONV_BASE: "blk.{bid}.ffn_conv_base", + MODEL_TENSOR.DFLASH_FFN_CONV_PROJ: "blk.{bid}.ffn_conv_proj", + MODEL_TENSOR.DFLASH_SELECTOR_PREV: "selector_predecessor", + MODEL_TENSOR.DFLASH_SELECTOR_NEXT: "selector_successor", + MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN: "selector_hidden", MODEL_TENSOR.D2T: "d2t", } @@ -2795,6 +2863,58 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ], + MODEL_ARCH.QWEN4EXP: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + # no OUTPUT_NORM / ATTN_NORM / ATTN_POST_NORM: hyper-connections replace every layer norm + MODEL_TENSOR.HC_HEAD_NORM, + MODEL_TENSOR.HC_HEAD_DOWN, + MODEL_TENSOR.HC_HEAD_UP, + MODEL_TENSOR.HC_ATTN_NORM, + MODEL_TENSOR.HC_ATTN_DOWN, + MODEL_TENSOR.HC_ATTN_UP, + MODEL_TENSOR.HC_ATTN_INJECT, + MODEL_TENSOR.HC_FFN_NORM, + MODEL_TENSOR.HC_FFN_DOWN, + MODEL_TENSOR.HC_FFN_UP, + MODEL_TENSOR.HC_FFN_INJECT, + # full attention layers: ATTN_Q holds [q|gate] interleaved per head + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.INDEXER_Q_PROJ, + MODEL_TENSOR.INDEXER_K_PROJ, + MODEL_TENSOR.INDEXER_Q_NORM, + MODEL_TENSOR.INDEXER_K_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_CONV1D, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_NORM, + MODEL_TENSOR.SSM_BETA, + MODEL_TENSOR.SSM_ALPHA, + MODEL_TENSOR.SSM_OUT, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_INP_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_GATE_UP_EXP, + MODEL_TENSOR.PER_LAYER_TOKEN_EMBD, + MODEL_TENSOR.PLE_KEY, + MODEL_TENSOR.PLE_VALUE, + MODEL_TENSOR.PLE_NORM_KEY, + MODEL_TENSOR.PLE_NORM_QUERY, + MODEL_TENSOR.PLE_NORM_CONV, + MODEL_TENSOR.PLE_CONV1D, + ], MODEL_ARCH.PLAMO: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, @@ -4953,6 +5073,13 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.DSPARK_MARKOV_W1, MODEL_TENSOR.DSPARK_MARKOV_W2, MODEL_TENSOR.DSPARK_CONF_PROJ, + MODEL_TENSOR.DFLASH_ATTN_CONV_BASE, + MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ, + MODEL_TENSOR.DFLASH_FFN_CONV_BASE, + MODEL_TENSOR.DFLASH_FFN_CONV_PROJ, + MODEL_TENSOR.DFLASH_SELECTOR_PREV, + MODEL_TENSOR.DFLASH_SELECTOR_NEXT, + MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN, ], MODEL_ARCH.MISTRAL4: [ MODEL_TENSOR.TOKEN_EMBD, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index d8a96a27bdd..d95fe9b1ac3 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -467,10 +467,15 @@ def write_tensors_to_file(self, *, progress: bool = False) -> None: shard_bar.reset(total=(total if total > 0 else None)) # relying on the fact that Python dicts preserve insertion order (since 3.7) - for ti in tensors.values(): + for name, ti in tensors.items(): assert ti.tensor is not None # can only iterate once over the tensors assert ti.tensor.nbytes == ti.nbytes + start = fout.tell() ti.tensor.tofile(fout) + # a short write here would only surface as a corrupt file at load time + if fout.tell() - start != ti.nbytes: + raise ValueError( + f"tensor {name!r} wrote {fout.tell() - start} bytes, expected {ti.nbytes}") if shard_bar is not None: shard_bar.update(ti.nbytes) if bar is not None: @@ -993,9 +998,24 @@ def add_sliding_window(self, value: int) -> None: def add_block_size(self, value: int) -> None: self.add_uint32(Keys.LLM.BLOCK_SIZE.format(arch=self.arch), value) + def add_conv_kernel_size(self, value: int) -> None: + self.add_uint32(Keys.LLM.CONV_KERNEL_SIZE.format(arch=self.arch), value) + + def add_conv_group_size(self, value: int) -> None: + self.add_uint32(Keys.LLM.CONV_GROUP_SIZE.format(arch=self.arch), value) + + def add_selector_rank(self, value: int) -> None: + self.add_uint32(Keys.LLM.SELECTOR_RANK.format(arch=self.arch), value) + + def add_selector_top_k(self, value: int) -> None: + self.add_uint32(Keys.LLM.SELECTOR_TOP_K.format(arch=self.arch), value) + def add_sample_from_anchor(self, value: bool) -> None: self.add_bool(Keys.LLM.SAMPLE_FROM_ANCHOR.format(arch=self.arch), value) + def add_has_confidence_head(self, value: bool) -> None: + self.add_bool(Keys.LLM.HAS_CONFIDENCE_HEAD.format(arch=self.arch), value) + def add_target_layers(self, value: Sequence[int]) -> None: self.add_array(Keys.LLM.TARGET_LAYERS.format(arch=self.arch), value) @@ -1029,6 +1049,40 @@ def add_hyper_connection_sinkhorn_iterations(self, count: int) -> None: def add_hyper_connection_epsilon(self, value: float) -> None: self.add_float32(Keys.HyperConnection.EPSILON.format(arch=self.arch), value) + def add_hyper_connection_low_rank(self, value: int) -> None: + self.add_uint32(Keys.HyperConnection.LOW_RANK.format(arch=self.arch), value) + + def add_ple_layers(self, values: Sequence[int]) -> None: + self.add_array(Keys.PerLayerEmbedding.LAYERS.format(arch=self.arch), values) + + def add_ple_ngram_size(self, value: int) -> None: + self.add_uint32(Keys.PerLayerEmbedding.NGRAM_SIZE.format(arch=self.arch), value) + + def add_ple_heads_per_ngram(self, value: int) -> None: + self.add_uint32(Keys.PerLayerEmbedding.HEADS_PER_NGRAM.format(arch=self.arch), value) + + def add_ple_conv_kernel(self, value: int) -> None: + self.add_uint32(Keys.PerLayerEmbedding.CONV_KERNEL.format(arch=self.arch), value) + + # multipliers reach ~2.4e13; default INT32 inference would truncate them + def _add_u64_array(self, key: str, values: Sequence[int]) -> None: + self.add_key_value(key, list(values), GGUFValueType.ARRAY, GGUFValueType.UINT64) + + def add_ple_layer_multipliers(self, values: Sequence[int]) -> None: + self._add_u64_array(Keys.PerLayerEmbedding.LAYER_MULTIPLIERS.format(arch=self.arch), values) + + def add_ple_head_offsets(self, values: Sequence[int]) -> None: + self._add_u64_array(Keys.PerLayerEmbedding.HEAD_OFFSETS.format(arch=self.arch), values) + + def add_ple_head_vocab_sizes(self, values: Sequence[int]) -> None: + self._add_u64_array(Keys.PerLayerEmbedding.HEAD_VOCAB_SIZES.format(arch=self.arch), values) + + def add_ple_eos_token_id(self, value: int) -> None: + self.add_uint32(Keys.PerLayerEmbedding.EOS_TOKEN_ID.format(arch=self.arch), value) + + def add_ple_image_token_id(self, value: int) -> None: + self.add_uint32(Keys.PerLayerEmbedding.IMAGE_TOKEN_ID.format(arch=self.arch), value) + def add_attention_scale(self, value: float) -> None: self.add_float32(Keys.Attention.SCALE.format(arch=self.arch), value) diff --git a/gguf-py/gguf/lazy.py b/gguf-py/gguf/lazy.py index acbc79258a3..a39f2232159 100644 --- a/gguf-py/gguf/lazy.py +++ b/gguf-py/gguf/lazy.py @@ -226,3 +226,68 @@ def tofile(self, *args, **kwargs): return eager.tofile(*args, **kwargs) # TODO: __array_function__ + + +# Tensor written to file one row-chunk at a time +class LazyChunkedTensor: + + def __init__( + self, chunks: list[Callable[[], np.ndarray]], shape: tuple[int, ...], dtype: DTypeLike, + qtype: Any = None, byteswap: bool = False, + ): + self._chunks = chunks + self._qtype = qtype + self._byteswap = byteswap + self.shape = tuple(shape) + self.dtype = np.dtype(dtype) + + @property + def nbytes(self) -> int: + n = self.dtype.itemsize + for d in self.shape: + n *= d + return n + + def numpy(self) -> LazyChunkedTensor: + return self + + def __array__(self, *args, **kwargs): + # numpy would otherwise make a 1-element object array of self, and write 8 bytes + raise TypeError("LazyChunkedTensor cannot become an ndarray, it is written in chunks") + + def quantize(self, qtype: Any) -> LazyChunkedTensor: + from .constants import GGMLQuantizationType + from .quants import QuantError, quant_shape_to_byte_shape + + if qtype == GGMLQuantizationType.F32: + shape, dtype = self.shape, np.dtype(np.float32) + elif qtype == GGMLQuantizationType.F16: + shape, dtype = self.shape, np.dtype(np.float16) + else: + try: + shape, dtype = quant_shape_to_byte_shape(self.shape, qtype), np.dtype(np.uint8) + except ValueError as e: + # raised here and not per chunk, so callers can still fall back to F16 + raise QuantError(str(e)) from e + return LazyChunkedTensor(self._chunks, shape, dtype, qtype, self._byteswap) + + def byteswap(self, inplace: bool = False) -> LazyChunkedTensor: + if inplace: + raise NotImplementedError("a chunked tensor cannot be byteswapped in place") + return LazyChunkedTensor(self._chunks, self.shape, self.dtype, self._qtype, not self._byteswap) + + def tofile(self, *args, **kwargs) -> None: + from .quants import quantize + + written = 0 + for load_chunk in self._chunks: + chunk = load_chunk() + if self._qtype is not None: + # exact only because chunks split on rows, and blocks never cross one + chunk = quantize(chunk, self._qtype) + if self._byteswap: + chunk = chunk.byteswap(inplace=False) + chunk.tofile(*args, **kwargs) + written += chunk.nbytes + del chunk + assert written == self.nbytes, f"chunked tensor wrote {written} bytes, expected {self.nbytes}" diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index ef580518e97..861acfe181f 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -1355,6 +1355,34 @@ class TensorNameMap: "model.confidence_head.proj", # dspark ), + MODEL_TENSOR.DFLASH_ATTN_CONV_BASE: ( + "model.layers.{bid}.attention_conv.base_kernel", + ), + + MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ: ( + "model.layers.{bid}.attention_conv.kernel_projection", + ), + + MODEL_TENSOR.DFLASH_FFN_CONV_BASE: ( + "model.layers.{bid}.mlp_conv.base_kernel", + ), + + MODEL_TENSOR.DFLASH_FFN_CONV_PROJ: ( + "model.layers.{bid}.mlp_conv.kernel_projection", + ), + + MODEL_TENSOR.DFLASH_SELECTOR_PREV: ( + "model.candidate_selector.predecessor_codebook", + ), + + MODEL_TENSOR.DFLASH_SELECTOR_NEXT: ( + "model.candidate_selector.successor_codebook", + ), + + MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN: ( + "model.candidate_selector.hidden_projection", + ), + MODEL_TENSOR.CLS: ( "classifier", # jina "classifier.dense", # roberta @@ -2680,6 +2708,65 @@ class TensorNameMap: "model.layers.{bid}.post_attention_layernorm", ), }, + MODEL_ARCH.QWEN4EXP: { + MODEL_TENSOR.HC_ATTN_NORM: ( + "model.layers.{bid}.attn_hyper_connection.hc_norm", + ), + MODEL_TENSOR.HC_ATTN_DOWN: ( + "model.layers.{bid}.attn_hyper_connection.input_mix_weight_down", + ), + MODEL_TENSOR.HC_ATTN_UP: ( + "model.layers.{bid}.attn_hyper_connection.input_mix_weight_up", + ), + MODEL_TENSOR.HC_ATTN_INJECT: ( + "model.layers.{bid}.attn_hyper_connection.block_inject_weight", + ), + MODEL_TENSOR.HC_FFN_NORM: ( + "model.layers.{bid}.mlp_hyper_connection.hc_norm", + ), + MODEL_TENSOR.HC_FFN_DOWN: ( + "model.layers.{bid}.mlp_hyper_connection.input_mix_weight_down", + ), + MODEL_TENSOR.HC_FFN_UP: ( + "model.layers.{bid}.mlp_hyper_connection.input_mix_weight_up", + ), + MODEL_TENSOR.HC_FFN_INJECT: ( + "model.layers.{bid}.mlp_hyper_connection.block_inject_weight", + ), + MODEL_TENSOR.HC_HEAD_NORM: ( + "model.hyper_connection_mixer.hc_norm", + ), + MODEL_TENSOR.HC_HEAD_DOWN: ( + "model.hyper_connection_mixer.input_mix_weight_down", + ), + MODEL_TENSOR.HC_HEAD_UP: ( + "model.hyper_connection_mixer.input_mix_weight_up", + ), + MODEL_TENSOR.INDEXER_Q_NORM: ( + "model.layers.{bid}.self_attn.indexer.q_layernorm", + ), + MODEL_TENSOR.INDEXER_K_NORM: ( + "model.layers.{bid}.self_attn.indexer.k_layernorm", + ), + MODEL_TENSOR.PLE_KEY: ( + "model.layers.{bid}.ple.key_proj", + ), + MODEL_TENSOR.PLE_VALUE: ( + "model.layers.{bid}.ple.value_proj", + ), + MODEL_TENSOR.PLE_NORM_KEY: ( + "model.layers.{bid}.ple.norm_key", + ), + MODEL_TENSOR.PLE_NORM_QUERY: ( + "model.layers.{bid}.ple.norm_query", + ), + MODEL_TENSOR.PLE_NORM_CONV: ( + "model.layers.{bid}.ple.norm_conv", + ), + MODEL_TENSOR.PLE_CONV1D: ( + "model.layers.{bid}.ple.conv1d", + ), + }, } mapping: dict[str, tuple[MODEL_TENSOR, str]] diff --git a/include/llama.h b/include/llama.h index a04177f9f7d..ef7a012c43a 100644 --- a/include/llama.h +++ b/include/llama.h @@ -43,10 +43,10 @@ #define LLAMA_FILE_MAGIC_GGSQ 0x67677371u // 'ggsq' #define LLAMA_SESSION_MAGIC LLAMA_FILE_MAGIC_GGSN -#define LLAMA_SESSION_VERSION 9 +#define LLAMA_SESSION_VERSION 10 #define LLAMA_STATE_SEQ_MAGIC LLAMA_FILE_MAGIC_GGSQ -#define LLAMA_STATE_SEQ_VERSION 2 +#define LLAMA_STATE_SEQ_VERSION 3 #ifdef __cplusplus extern "C" { @@ -214,6 +214,12 @@ extern "C" { LLAMA_API const char * llama_load_mode_name(enum llama_load_mode load_mode); LLAMA_API enum llama_load_mode llama_load_mode_from_str(const char * str); + enum llama_lazy_mode { + LLAMA_LAZY_MODE_OFF = 0, // always read the whole tensor up front + LLAMA_LAZY_MODE_AUTO = 1, // lazy only for marked tensors larger than 4 GiB (requires mmap) + LLAMA_LAZY_MODE_ON = 2, // read the rows of tensors marked by the arch on demand (requires mmap) + }; + enum llama_context_type { LLAMA_CONTEXT_TYPE_DEFAULT = 0, LLAMA_CONTEXT_TYPE_MTP = 1, @@ -315,6 +321,8 @@ extern "C" { enum llama_split_mode split_mode; // how to split the model across multiple GPUs enum llama_load_mode load_mode; // how to load the model + enum llama_lazy_mode lazy_mode; // on-demand reading of tensors marked by the arch + // the GPU that is used for the entire model when split_mode is LLAMA_SPLIT_MODE_NONE int32_t main_gpu; @@ -437,6 +445,7 @@ extern "C" { const struct llama_model_kv_override * kv_overrides; // pointer to kv overrides const struct llama_model_tensor_override * tt_overrides; // pointer to tensor overrides const int32_t * prune_layers; // pointer to layer indices to prune + size_t max_buf_size; // max bytes of tensor rows kept in memory at once, 0 = default (8 GiB) } llama_model_quantize_params; typedef struct llama_logit_bias { diff --git a/scripts/pr2wt.sh b/scripts/pr2wt.sh index ae03a888cb6..ce0327eadb5 100755 --- a/scripts/pr2wt.sh +++ b/scripts/pr2wt.sh @@ -48,7 +48,11 @@ echo "org/repo: $org_repo" meta=$(curl -sSLf -H "Accept: application/vnd.github+json" "https://api.github.com/repos/$org_repo/pulls/$PR") -url_remote=$(echo "$meta" | jq -r '.head.repo.clone_url') +if [[ $url_origin =~ ^git@ ]]; then + url_remote=$(echo "$meta" | jq -r '.head.repo.ssh_url') +else + url_remote=$(echo "$meta" | jq -r '.head.repo.clone_url') +fi head_ref=$(echo "$meta" | jq -r '.head.ref') echo "url: $url_remote" diff --git a/scripts/snapdragon/adb/run-bench.sh b/scripts/snapdragon/adb/run-bench.sh deleted file mode 100755 index eaae80a77d6..00000000000 --- a/scripts/snapdragon/adb/run-bench.sh +++ /dev/null @@ -1,49 +0,0 @@ -#!/bin/sh -# - -# Basedir on device -basedir=/data/local/tmp/llama.cpp - -branch=. -[ "$B" != "" ] && branch=$B - -adbserial= -[ "$S" != "" ] && adbserial="-s $S" - -adbhost= -[ "$H" != "" ] && adbhost="-H $H" - -model="Llama-3.2-3B-Instruct-Q4_0.gguf" -[ "$M" != "" ] && model="$M" - -device="HTP0" -[ "$D" != "" ] && device="$D" - -verbose= -[ "$V" != "" ] && verbose="GGML_HEXAGON_VERBOSE=$V" cli_opts="$cli_opts -v" - -profile= -[ "$PROF" != "" ] && profile="GGML_HEXAGON_PROFILE=$PROF" cli_opts="$cli_opts -v" - -opmask= -[ "$OPSTAGE" != "" ] && opmask="GGML_HEXAGON_OPSTAGE=$OPSTAGE" - -nhvx= -[ "$NHVX" != "" ] && nhvx="GGML_HEXAGON_NHVX=$NHVX" - -ndev= -[ "$NDEV" != "" ] && ndev="GGML_HEXAGON_NDEV=$NDEV" - -hb= -[ "$HB" != "" ] && hb="GGML_HEXAGON_HOSTBUF=$HB" - -set -x - -adb $adbserial $adbhost shell " \ - cd $basedir; \ - LD_LIBRARY_PATH=$basedir/$branch/lib \ - ADSP_LIBRARY_PATH=$basedir/$branch/lib \ - $ndev $nhvx $opmask $verbose $profile $hb ./$branch/bin/llama-bench --device $device --load-mode none -m $basedir/../gguf/$model \ - --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 \ - --ubatch-size 1024 -fa 1 -ngl 99 $cli_opts $@ \ -" diff --git a/scripts/snapdragon/adb/run-cli.sh b/scripts/snapdragon/adb/run-cli.sh deleted file mode 100755 index 27a4a14195f..00000000000 --- a/scripts/snapdragon/adb/run-cli.sh +++ /dev/null @@ -1,78 +0,0 @@ -#!/bin/sh -# - -# Basedir on device -basedir=/data/local/tmp/llama.cpp - -cli_opts= - -branch=. -[ "$B" != "" ] && branch=$B - -adbserial= -[ "$S" != "" ] && adbserial="-s $S" - -adbhost= -[ "$H" != "" ] && adbhost="-H $H" - -model="Llama-3.2-3B-Instruct-Q4_0.gguf" -[ "$M" != "" ] && model="$M" - -device="HTP0" -[ "$D" != "" ] && device="$D" - -verbose= -[ "$V" != "" ] && verbose="GGML_HEXAGON_VERBOSE=$V" cli_opts="$cli_opts -v" - -sched= -[ "$SCHED" != "" ] && sched="GGML_SCHED_DEBUG=2" cli_opts="$cli_opts -v" - -profile= -[ "$PROF" != "" ] && profile="GGML_HEXAGON_PROFILE=$PROF" cli_opts="$cli_opts -v" - -opmask= -[ "$OPSTAGE" != "" ] && opmask="GGML_HEXAGON_OPSTAGE=$OPSTAGE" - -nhvx= -[ "$NHVX" != "" ] && nhvx="GGML_HEXAGON_NHVX=$NHVX" - -hmx= -[ "$HMX" != "" ] && hmx="GGML_HEXAGON_USE_HMX=$HMX" - -ndev= -[ "$NDEV" != "" ] && ndev="GGML_HEXAGON_NDEV=$NDEV" - -hb= -[ "$HB" != "" ] && hb="GGML_HEXAGON_HOSTBUF=$HB" - -opbatch= -[ "$OB" != "" ] && opbatch="GGML_HEXAGON_OPBATCH=$OB" - -opqueue= -[ "$OQ" != "" ] && opqueue="GGML_HEXAGON_OPQUEUE=$OQ" - -opflt= -[ "$OF" != "" ] && opflt="GGML_HEXAGON_OPFILTER=$OF" - -vmem= -[ "$VM" != "" ] && opflt="GGML_HEXAGON_VMEM=$VM" - -mbuf= -[ "$MB" != "" ] && opflt="GGML_HEXAGON_MBUF=$MB" -vmem= -[ "$VM" != "" ] && vmem="GGML_HEXAGON_VMEM=$VM" - -mbuf= -[ "$MB" != "" ] && mbuf="GGML_HEXAGON_MBUF=$MB" -set -x - -adb $adbserial $adbhost shell " \ - cd $basedir; ulimit -c unlimited; \ - LD_LIBRARY_PATH=$basedir/$branch/lib \ - ADSP_LIBRARY_PATH=$basedir/$branch/lib \ - $verbose $sched $opmask $profile $nhvx $hmx $ndev $hb $opbatch $opqueue $opflt $vmem $mbuf \ - ./$branch/bin/llama-cli --load-mode none -m $basedir/../gguf/$model \ - --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 \ - --ctx-size 8192 --ubatch-size 1024 -fa on \ - -ngl 99 --device $device $cli_opts $@ \ -" diff --git a/scripts/snapdragon/adb/run-completion.sh b/scripts/snapdragon/adb/run-completion.sh deleted file mode 100755 index 30893ed293a..00000000000 --- a/scripts/snapdragon/adb/run-completion.sh +++ /dev/null @@ -1,86 +0,0 @@ -#!/bin/sh -# - -# Basedir on device -basedir=/data/local/tmp/llama.cpp - -cli_opts= - -branch=. -[ "$B" != "" ] && branch=$B - -adbserial= -[ "$S" != "" ] && adbserial="-s $S" - -adbhost= -[ "$H" != "" ] && adbhost="-H $H" - -model="Llama-3.2-3B-Instruct-Q4_0.gguf" -[ "$M" != "" ] && model="$M" - -device="HTP0" -[ "$D" != "" ] && device="$D" - -verbose= -[ "$V" != "" ] && verbose="GGML_HEXAGON_VERBOSE=$V" cli_opts="$cli_opts -v" - -sched= -[ "$SCHED" != "" ] && sched="GGML_SCHED_DEBUG=2" cli_opts="$cli_opts -v" - -profile= -[ "$PROF" != "" ] && profile="GGML_HEXAGON_PROFILE=$PROF" cli_opts="$cli_opts -v" - -opmask= -[ "$OPSTAGE" != "" ] && opmask="GGML_HEXAGON_OPSTAGE=$OPSTAGE" - -nhvx= -[ "$NHVX" != "" ] && nhvx="GGML_HEXAGON_NHVX=$NHVX" - -hmx= -[ "$HMX" != "" ] && hmx="GGML_HEXAGON_USE_HMX=$HMX" - -ndev= -[ "$NDEV" != "" ] && ndev="GGML_HEXAGON_NDEV=$NDEV" - -hb= -[ "$HB" != "" ] && hb="GGML_HEXAGON_HOSTBUF=$HB" - -opbatch= -[ "$OB" != "" ] && opbatch="GGML_HEXAGON_OPBATCH=$OB" - -opqueue= -[ "$OQ" != "" ] && opqueue="GGML_HEXAGON_OPQUEUE=$OQ" - -oppoll= -[ "$OP" != "" ] && oppoll="GGML_HEXAGON_OPPOLL=$OP" - -opflt= -[ "$OF" != "" ] && opflt="GGML_HEXAGON_OPFILTER=$OF" - -opfuse= -[ "$OC" != "" ] && opfuse="GGML_HEXAGON_OPFUSION=$OC" - -vmem= -[ "$VM" != "" ] && vmem="GGML_HEXAGON_VMEM=$VM" - -mbuf= -[ "$MB" != "" ] && mbuf="GGML_HEXAGON_MBUF=$MB" - -mmsel= -[ "$MM" != "" ] && mmsel="GGML_HEXAGON_MM_SELECT=$MM" - -fasel= -[ "$FA" != "" ] && fasel="GGML_HEXAGON_FA_SELECT=$FA" - -set -x - -adb $adbserial $adbhost shell " \ - cd $basedir; ulimit -c unlimited; \ - LD_LIBRARY_PATH=$basedir/$branch/lib \ - ADSP_LIBRARY_PATH=$basedir/$branch/lib \ - $verbose $sched $opmask $profile $nhvx $hmx $ndev $hb $opbatch $opqueue $oppoll $opflt $opfuse $vmem $mbuf $mmsel $fasel \ - ./$branch/bin/llama-completion --load-mode none -m $basedir/../gguf/$model \ - --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 \ - --ctx-size 8192 --ubatch-size 1024 -fa on \ - -ngl 99 --device $device $cli_opts $@ \ -" diff --git a/scripts/snapdragon/adb/run-mtmd.sh b/scripts/snapdragon/adb/run-mtmd.sh deleted file mode 100755 index 65dd6ec59e5..00000000000 --- a/scripts/snapdragon/adb/run-mtmd.sh +++ /dev/null @@ -1,71 +0,0 @@ -#!/bin/sh -# - -# Basedir on device -basedir=/data/local/tmp/llama.cpp - -cli_opts= - -branch=. -[ "$B" != "" ] && branch=$B - -adbserial= -[ "$S" != "" ] && adbserial="-s $S" - -adbhost= -[ "$H" != "" ] && adbhost="-H $H" - -model="gemma-3-4b-it-Q4_0.gguf" -[ "$M" != "" ] && model="$M" - -mmproj="mmproj-F16.gguf" -[ "$MMPROJ" != "" ] && mmproj="$MMPROJ" - -image= -[ "$IMG" != "" ] && image="$IMG" - -device="HTP0" -[ "$D" != "" ] && device="$D" - -verbose= -[ "$V" != "" ] && verbose="GGML_HEXAGON_VERBOSE=$V" - -experimental="GGML_HEXAGON_EXPERIMENTAL=1" -[ "$E" != "" ] && experimental="GGML_HEXAGON_EXPERIMENTAL=$E" - -sched= -[ "$SCHED" != "" ] && sched="GGML_SCHED_DEBUG=2" cli_opts="$cli_opts -v" - -profile= -[ "$PROF" != "" ] && profile="GGML_HEXAGON_PROFILE=$PROF" - -opmask= -[ "$OPSTAGE" != "" ] && opmask="GGML_HEXAGON_OPSTAGE=$OPSTAGE" - -nhvx= -[ "$NHVX" != "" ] && nhvx="GGML_HEXAGON_NHVX=$NHVX" - -hmx= -[ "$HMX" != "" ] && hmx="GGML_HEXAGON_USE_HMX=$HMX" - -ndev= -[ "$NDEV" != "" ] && ndev="GGML_HEXAGON_NDEV=$NDEV" - -# MTMD backend device for vision model (defaults to CPU if not set) -mtmd_backend= -[ "$MTMD_DEVICE" != "" ] && mtmd_backend="MTMD_BACKEND_DEVICE=$MTMD_DEVICE" - -set -x - -adb $adbserial $adbhost shell " \ - cd $basedir; ulimit -c unlimited; \ - LD_LIBRARY_PATH=$basedir/$branch/lib \ - ADSP_LIBRARY_PATH=$basedir/$branch/lib \ - $verbose $experimental $sched $opmask $profile $hmx $nhvx $ndev $mtmd_backend \ - ./$branch/bin/llama-mtmd-cli --load-mode none -m $basedir/../gguf/$model \ - --mmproj $basedir/../gguf/$mmproj \ - --image $basedir/../gguf/$image \ - --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 \ - --ctx-size 8192 --ubatch-size 1024 -fa on \ - -ngl 99 --device $device -v $cli_opts $@ \ -" diff --git a/scripts/snapdragon/adb/run-tool.sh b/scripts/snapdragon/adb/run-tool.sh deleted file mode 100755 index 8bf7ba3a546..00000000000 --- a/scripts/snapdragon/adb/run-tool.sh +++ /dev/null @@ -1,72 +0,0 @@ -#!/bin/sh -# - -# Basedir on device -basedir=/data/local/tmp/llama.cpp - -cli_opts= - -branch=. -[ "$B" != "" ] && branch=$B - -adbserial= -[ "$S" != "" ] && adbserial="-s $S" - -adbhost= -[ "$H" != "" ] && adbhost="-H $H" - -device="HTP0" -[ "$D" != "" ] && device="$D" - -verbose= -[ "$V" != "" ] && verbose="GGML_HEXAGON_VERBOSE=$V" - -sched= -[ "$SCHED" != "" ] && sched="GGML_SCHED_DEBUG=2" cli_opts="$cli_opts -v" - -profile= -[ "$PROF" != "" ] && profile="GGML_HEXAGON_PROFILE=$PROF" - -opmask= -[ "$OPSTAGE" != "" ] && opmask="GGML_HEXAGON_OPSTAGE=$OPSTAGE" - -nhvx= -[ "$NHVX" != "" ] && nhvx="GGML_HEXAGON_NHVX=$NHVX" - -hmx= -[ "$HMX" != "" ] && hmx="GGML_HEXAGON_USE_HMX=$HMX" - -ndev= -[ "$NDEV" != "" ] && ndev="GGML_HEXAGON_NDEV=$NDEV" - -hb= -[ "$HB" != "" ] && hb="GGML_HEXAGON_HOSTBUF=$HB" - -opbatch= -[ "$OB" != "" ] && opbatch="GGML_HEXAGON_OPBATCH=$OB" - -opqueue= -[ "$OQ" != "" ] && opqueue="GGML_HEXAGON_OPQUEUE=$OQ" - -oppoll= -[ "$OP" != "" ] && oppoll="GGML_HEXAGON_OPPOLL=$OP" - -opfuse= -[ "$OC" != "" ] && opfuse="GGML_HEXAGON_OPFUSION=$OC" - -mmsel= -[ "$MM" != "" ] && mmsel="GGML_HEXAGON_MM_SELECT=$MM" - -fasel= -[ "$FA" != "" ] && fasel="GGML_HEXAGON_FA_SELECT=$FA" - -set -x - -tool=$1; shift - -adb $adbserial $adbhost shell " \ - cd $basedir; ulimit -c unlimited; \ - LD_LIBRARY_PATH=$basedir/$branch/lib \ - ADSP_LIBRARY_PATH=$basedir/$branch/lib \ - $verbose $sched $opmask $profile $nhvx $hmx $ndev $hb $opbatch $opqueue $oppoll $opfuse $mmsel $fasel ./$branch/bin/$tool $@ \ -" diff --git a/scripts/snapdragon/build.py b/scripts/snapdragon/build.py new file mode 100755 index 00000000000..5e9fab3d08b --- /dev/null +++ b/scripts/snapdragon/build.py @@ -0,0 +1,266 @@ +#!/usr/bin/env python3 +# +# Build llama.cpp for Snapdragon (via Docker or natively) and push to device. +# + +import sys +import os +import argparse +import subprocess +import platform +import shutil +import logging + +from sdk import validate_windows_sdks + +logger = logging.getLogger("build") + + +def parse_target(target_str): + if not target_str: + return None, None + if target_str.startswith("adb") or target_str.startswith("android"): + parts = target_str.split(":", 1) + serial = parts[1] if len(parts) > 1 else None + return "android", serial + elif target_str.startswith("lnx") or target_str.startswith("linux") or target_str.startswith("ubuntu"): + parts = target_str.split(":", 1) + host = parts[1] if len(parts) > 1 else None + return "linux", host + elif target_str in ("wos", "windows"): + return "windows", None + else: + return None, None + + +def get_uid_gid(): + if platform.system() != "Windows": + return [f"{os.getuid()}:{os.getgid()}"] + return [] + + +def main(): + logging.basicConfig(level=logging.INFO, format='%(message)s') + parser = argparse.ArgumentParser( + description="Build llama.cpp for Snapdragon using cross-compilation docker containers or natively." + ) + parser.add_argument("--target", default="android", help="Compilation target and deployment definition (e.g. android[:serial]/adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, windows/wos) (default: android)") + parser.add_argument("--build-dir", help="Build directory name (defaults to build-TARGET[-dbg], e.g. build-android)") + parser.add_argument("--install-dir", help="Install directory name (defaults to pkg-TARGET[-dbg], e.g. pkg-android)") + parser.add_argument("--jobs", "-j", type=int, help="Number of build jobs (defaults to CPU thread count)") + parser.add_argument("--no-docker", action="store_true", help="Build natively on the host instead of in a docker container") + parser.add_argument("--preset", help="Override the CMake preset to use") + parser.add_argument("--debug", action="store_true", help="Build in debug mode (uses -debug presets instead of -release)") + + # Push options + parser.add_argument("--push", action="store_true", help="Push built package to the target device via ADB or SSH/SCP") + parser.add_argument("--target-dir", help="Target directory on the device (default: /data/local/tmp/llama.cpp for Android, ~/llama.cpp for Linux)") + + # Toolchain options + parser.add_argument("--toolchain-version", default="v0.7", help="Docker toolchain image version/tag (default: v0.7)") + parser.add_argument("--toolchain-url", default="ghcr.io/snapdragon-toolchain", help="Docker toolchain registry URL/namespace (default: ghcr.io/snapdragon-toolchain)") + + args = parser.parse_args() + + target_type, target_val = parse_target(args.target) + if not target_type: + logger.error(f"Error: Invalid target format '{args.target}'. Must be android[:serial]/adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, or windows/wos.") + sys.exit(1) + + if target_type == "windows": + logger.info("Windows target selected. Forcing native compilation...") + args.no_docker = True + if platform.system() != "Windows": + logger.warning("Warning: Windows compilation is intended to run on Windows arm64 hosts.") + validate_windows_sdks() + + # Determine preset and check if it's debug + preset = args.preset + if preset: + is_debug = args.debug or ("debug" in preset.lower()) + else: + is_debug = args.debug + config_type = "debug" if is_debug else "release" + if args.no_docker: + if target_type == "windows" or platform.system() == "Windows": + preset = f"arm64-windows-snapdragon-{config_type}" + elif target_type == "linux": + preset = f"arm64-linux-snapdragon-{config_type}" + else: + preset = f"arm64-android-snapdragon-{config_type}" + else: + preset = f"arm64-linux-snapdragon-{config_type}" if target_type == "linux" else f"arm64-android-snapdragon-{config_type}" + + target_prefix = args.target.split(":", 1)[0] + suffix = "-dbg" if is_debug else "" + + build_dir = args.build_dir + if not build_dir: + build_dir = f"build-{target_prefix}{suffix}" + + install_dir = args.install_dir + if not install_dir: + install_dir = f"pkg-{target_prefix}{suffix}" + + repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) + + # Ensure CMakeUserPresets.json is in the workspace root, update if docs version is newer + preset_src = os.path.join(repo_root, "docs", "backend", "snapdragon", "CMakeUserPresets.json") + preset_dst = os.path.join(repo_root, "CMakeUserPresets.json") + if os.path.exists(preset_src): + should_copy = False + if not os.path.exists(preset_dst): + should_copy = True + else: + # Check modification times + src_mtime = os.path.getmtime(preset_src) + dst_mtime = os.path.getmtime(preset_dst) + if src_mtime > dst_mtime: + preset_bak = preset_dst + ".bak" + logger.info(f"Docs CMakeUserPresets.json is newer. Backing up existing {preset_dst} to {preset_bak}") + shutil.copy2(preset_dst, preset_bak) + should_copy = True + + if should_copy: + logger.info(f"Copying CMakeUserPresets.json from {preset_src} to {preset_dst}") + shutil.copy2(preset_src, preset_dst) + else: + logger.warning("Warning: CMakeUserPresets.json not found in docs/backend/snapdragon/.") + + jobs = args.jobs if args.jobs else os.cpu_count() or 4 + + if args.no_docker: + # Native/local host build + logger.info("Running native/local CMake build...") + install_prefix = os.path.join(repo_root, install_dir, "llama.cpp") + + # Configure + configure_cmd = ["cmake", f"--preset={preset}", "-B", build_dir] + logger.info(f"+ {' '.join(configure_cmd)}") + res = subprocess.run(configure_cmd, cwd=repo_root) + if res.returncode != 0: + logger.error("CMake configuration failed.") + sys.exit(res.returncode) + + # Build + build_cmd = ["cmake", "--build", build_dir, "-j", str(jobs)] + logger.info(f"+ {' '.join(build_cmd)}") + res = subprocess.run(build_cmd, cwd=repo_root) + if res.returncode != 0: + logger.error("CMake build failed.") + sys.exit(res.returncode) + + # Install + install_cmd = ["cmake", "--install", build_dir, "--prefix", install_prefix] + logger.info(f"+ {' '.join(install_cmd)}") + res = subprocess.run(install_cmd, cwd=repo_root) + if res.returncode != 0: + logger.error("CMake install failed.") + sys.exit(res.returncode) + else: + # Docker-based build + logger.info("Running Docker-based cross-compilation build...") + image_name = "arm64-linux" if target_type == "linux" else "arm64-android" + image = f"{args.toolchain_url}/{image_name}:{args.toolchain_version}" + + install_prefix_container = f"/workspace/{install_dir}/llama.cpp" + + build_sh_cmd = ( + f"cmake --preset {preset} -B /workspace/{build_dir} && " + f"cmake --build /workspace/{build_dir} -j {jobs} && " + f"cmake --install /workspace/{build_dir} --prefix {install_prefix_container}" + ) + + docker_cmd = [ + "docker", "run", "--rm", + "--volume", f"{repo_root}:/workspace", + "--workdir", "/workspace", + "--platform", "linux/amd64" + ] + uid_gid = get_uid_gid() + if uid_gid: + docker_cmd += ["-u", uid_gid[0]] + + docker_cmd += [image, "bash", "-c", build_sh_cmd] + + logger.info(f"+ {' '.join(docker_cmd)}") + res = subprocess.run(docker_cmd, cwd=repo_root) + if res.returncode != 0: + logger.error("Docker-based build failed.") + sys.exit(res.returncode) + + logger.info("\nBuild and installation completed successfully!") + + # Push/deploy if requested + if args.push: + src_path = os.path.join(repo_root, install_dir, "llama.cpp") + if not os.path.exists(src_path): + logger.error(f"Error: installation directory {src_path} does not exist. Cannot deploy.") + sys.exit(1) + + # Resolve target directory on device + target_dir = args.target_dir + if not target_dir: + target_dir = "/data/local/tmp/llama.cpp" if target_type == "android" else "~/llama.cpp" + target_dir = target_dir.rstrip("/") + + sub_items = [item for item in os.listdir(src_path) if not item.startswith(".")] + + if target_type == "android": + logger.info("\nPushing built artifacts to Android device via ADB...") + adb_cmd = ["adb"] + if target_val: # serial + adb_cmd += ["-s", target_val] + + # Clean stale package files on device + if sub_items: + clean_paths = " ".join(f"{target_dir}/{item}" for item in sub_items) + clean_cmd = adb_cmd + ["shell", f"rm -rf {clean_paths}"] + logger.info(f"+ {' '.join(clean_cmd)}") + subprocess.run(clean_cmd) + + # Android destination directory is target_dir + push_cmd = adb_cmd + ["push", os.path.join(src_path, "."), target_dir] + logger.info(f"+ {' '.join(push_cmd)}") + res = subprocess.run(push_cmd) + if res.returncode != 0: + logger.error("ADB push failed.") + sys.exit(res.returncode) + logger.info("ADB push completed successfully!") + + elif target_type == "linux": + ssh_host = target_val + if not ssh_host: + logger.error("Error: SSH host not specified in target (e.g. use linux:user@host, lnx:user@host, or ubuntu:user@host). Cannot deploy.") + sys.exit(1) + logger.info(f"\nDeploying built artifacts to Linux device {ssh_host} via SSH/SCP...") + + # Clean stale package files on remote host + if sub_items: + clean_paths = " ".join(f"{target_dir}/{item}" for item in sub_items) + clean_cmd = ["ssh", ssh_host, f"rm -rf {clean_paths}"] + logger.info(f"+ {' '.join(clean_cmd)}") + subprocess.run(clean_cmd) + + # Deploy to target_dir + deploy_cmd = ["scp", "-r", os.path.join(src_path, "."), f"{ssh_host}:{target_dir}"] + logger.info(f"+ {' '.join(deploy_cmd)}") + res = subprocess.run(deploy_cmd) + if res.returncode != 0: + logger.error("SSH/SCP deploy failed.") + sys.exit(res.returncode) + logger.info("SSH/SCP deploy completed successfully!") + + elif target_type == "windows": + logger.info("\nPush for Windows on Snapdragon (windows) target is currently a stub.") + + +if __name__ == "__main__": + try: + main() + except KeyboardInterrupt: + logger.info("\nInterrupted by user.") + sys.exit(130) + except RuntimeError as err: + logger.error("Error: %s", err) + sys.exit(1) diff --git a/scripts/snapdragon/ggml-hexagon-profile.py b/scripts/snapdragon/ggml-hexagon-profile.py index 97a3acd26c2..038d92fb5c4 100755 --- a/scripts/snapdragon/ggml-hexagon-profile.py +++ b/scripts/snapdragon/ggml-hexagon-profile.py @@ -34,6 +34,26 @@ r"trace-evt\s+(?P[A-Z_0-9\-]+):\s+thread\s+(?P\d+)\s+info\s+(?P\d+)\s+(?Pstart|stop)\s+(?P\d+)" ) +device_pattern = re.compile(r"\b(HTP\d+(?::\d+)?)\s+(?:profile-op|trace-evt)\b") + + +def extract_device(line): + m = device_pattern.search(line) + if m: + return m.group(1) + return "HTP0" + + +def device_matches(record_device, target_device): + targets = [t.strip() for t in target_device.split(',')] + for target in targets: + if record_device == target: + return True + if record_device.startswith(target + ":"): + return True + return False + + logger = logging.getLogger("ggml-hexagon-profile") @@ -72,7 +92,7 @@ def unwrap(self, raw): return raw + self.high_part -def parse_log(file_path, pmu_index=None): +def parse_log(file_path, pmu_index=None, limit=None, device_filter=None, op_filter_re=None): try: if file_path != "-": f = open(file_path, 'r', encoding='utf-8', errors='ignore') @@ -85,13 +105,22 @@ def parse_log(file_path, pmu_index=None): all_ops: List[Dict[str, Any]] = [] all_traces: List[Dict[str, Any]] = [] current_op: Optional[Dict[str, Any]] = None + ops_count_per_device = {} + if device_filter is not None: + for target in device_filter.split(','): + ops_count_per_device[target.strip()] = 0 + limit_reached = False - timestamp_pattern = re.compile(r"^(?P\d+)\.(?P\d+)\.(?P\d+)\.(?P\d+)\s+[A-Z]\s+") - unwrapper = None - trace_unwrapper = None + timestamp_pattern = re.compile(r"(?P\d+)\.(?P\d+)\.(?P\d+)\.(?P\d+)\s+[A-Z]\s+") + unwrappers = {} + last_batch_start = {} + trace_unwrappers = {} for line in f: - ts_match = timestamp_pattern.match(line) + if "profile-op" not in line and "trace-evt" not in line: + continue + + ts_match = timestamp_pattern.search(line) abs_usec = 0 if ts_match: abs_usec = ( @@ -100,8 +129,11 @@ def parse_log(file_path, pmu_index=None): + int(ts_match.group('us')) ) - if "|" in line and "profile-op" in line: - parts = [p.strip() for p in line.split("|")] + device = extract_device(line) + + idx = line.find("profile-op") + if idx != -1 and "|" in line[idx:]: + parts = [p.strip() for p in line[idx:].split("|")] prefix = parts[0] prefix_match = re.search(r"profile-op\s+(?P[A-Z_0-9+]+)", prefix) if not prefix_match: @@ -145,7 +177,6 @@ def parse_log(file_path, pmu_index=None): except (ValueError, IndexError): pmu_val = None - evt_val = None evt_val = None if types.startswith("evt-cnt "): try: @@ -158,14 +189,18 @@ def parse_log(file_path, pmu_index=None): if op_name == "OPBATCH": if cycles_start_raw: unwrapped_cycles_start = int(cycles_start_raw) - unwrapper = CycleUnwrapper(unwrapped_cycles_start) - trace_unwrapper = CycleUnwrapper(unwrapped_cycles_start) + unwrappers[device] = CycleUnwrapper(unwrapped_cycles_start) + last_batch_start[device] = unwrapped_cycles_start + for k in list(trace_unwrappers.keys()): + if k[0] == device: + del trace_unwrappers[k] else: - if cycles_start_raw and unwrapper is not None: - unwrapped_cycles_start = unwrapper.unwrap(int(cycles_start_raw)) + if cycles_start_raw: + device_unwrapper = unwrappers.get(device) + if device_unwrapper is not None: + unwrapped_cycles_start = device_unwrapper.unwrap(int(cycles_start_raw)) - idx = line.find("profile-op ") - op_text = line[idx + 11:].strip() if idx != -1 else line.strip() + op_text = re.sub(r"^profile-op\s+", "", line[idx:]).strip() if idx != -1 else line.strip() current_op = { 'name': op_name, @@ -180,24 +215,58 @@ def parse_log(file_path, pmu_index=None): 'pmu_val': pmu_val, 'evt_val': evt_val, 'abs_usec': abs_usec, - 'trace_events': [] + 'trace_events': [], + 'device': device } all_ops.append(current_op) + + # Check if matching early exit criteria + matched = False + matched_target = None + if device_filter is not None: + targets = [t.strip() for t in device_filter.split(',')] + for target in targets: + if device == target or device.startswith(target + ":"): + matched = True + matched_target = target + break + else: + matched = True + matched_target = device + + if op_filter_re is not None and not op_filter_re.search(op_text): + matched = False + + if matched: + if matched_target not in ops_count_per_device: + ops_count_per_device[matched_target] = 0 + ops_count_per_device[matched_target] += 1 + + if limit is not None and len(ops_count_per_device) > 0 and all(count >= limit for count in ops_count_per_device.values()): + limit_reached = True + + if limit_reached and op_name == "OPBATCH": + break continue trace_match = trace_pattern.search(line) if trace_match: + thread = int(trace_match.group('thread')) raw_cyc = int(trace_match.group('cycles')) unwrapped_cyc = None - if trace_unwrapper is not None: - unwrapped_cyc = trace_unwrapper.unwrap(raw_cyc) + th_key = (device, thread) + if th_key not in trace_unwrappers: + batch_start = last_batch_start.get(device) + trace_unwrappers[th_key] = CycleUnwrapper(batch_start) + unwrapped_cyc = trace_unwrappers[th_key].unwrap(raw_cyc) all_traces.append({ - 'thread': int(trace_match.group('thread')), + 'thread': thread, 'event': trace_match.group('event'), 'info': int(trace_match.group('info')), 'cycles': raw_cyc, 'unwrapped_cycles': unwrapped_cyc, - 'state': trace_match.group('state') + 'state': trace_match.group('state'), + 'device': device }) f.close() @@ -207,39 +276,45 @@ def parse_log(file_path, pmu_index=None): op['start_cycles'] = op['unwrapped_cycles_start'] op['end_cycles'] = op['start_cycles'] + op['cycles'] if op['start_cycles'] is not None else None - # Filter ops with valid start_cycles - valid_ops = [op for op in all_ops if op['start_cycles'] is not None and op['end_cycles'] is not None] + # Group ops by device + valid_ops_by_dev = defaultdict(list) + for op in all_ops: + if op['start_cycles'] is not None and op['end_cycles'] is not None: + valid_ops_by_dev[op['device']].append(op) + + # Group trace events by device + traces_by_dev = defaultdict(list) + for e in all_traces: + if e['unwrapped_cycles'] is not None: + traces_by_dev[e['device']].append(e) - # Separate OPBATCH ops from other ops - opbatch_ops = [op for op in valid_ops if op['name'] == "OPBATCH"] - other_ops = [op for op in valid_ops if op['name'] != "OPBATCH"] + for device, dev_ops in valid_ops_by_dev.items(): + opbatch_ops = [op for op in dev_ops if op['name'] == "OPBATCH"] + other_ops = [op for op in dev_ops if op['name'] != "OPBATCH"] - # Sort them by start_cycles to enable binary search - opbatch_ops.sort(key=lambda op: op['start_cycles']) - other_ops.sort(key=lambda op: op['start_cycles']) + opbatch_ops.sort(key=lambda op: op['start_cycles']) + other_ops.sort(key=lambda op: op['start_cycles']) - opbatch_starts = [op['start_cycles'] for op in opbatch_ops] - other_starts = [op['start_cycles'] for op in other_ops] + opbatch_starts = [op['start_cycles'] for op in opbatch_ops] + other_starts = [op['start_cycles'] for op in other_ops] - # Map trace events to any operator whose cycles contain them - for e in all_traces: - cyc = e['unwrapped_cycles'] - if cyc is None: - continue + dev_traces = traces_by_dev.get(device, []) + for e in dev_traces: + cyc = e['unwrapped_cycles'] - # Map to OPBATCH - idx = bisect.bisect_right(opbatch_starts, cyc) - 1 - if idx >= 0: - op = opbatch_ops[idx] - if op['start_cycles'] <= cyc <= op['end_cycles']: - op['trace_events'].append(e) + # Map to OPBATCH + idx = bisect.bisect_right(opbatch_starts, cyc) - 1 + if idx >= 0: + op = opbatch_ops[idx] + if op['start_cycles'] <= cyc <= op['end_cycles']: + op['trace_events'].append(e) - # Map to other ops - idx = bisect.bisect_right(other_starts, cyc) - 1 - if idx >= 0: - op = other_ops[idx] - if op['start_cycles'] <= cyc <= op['end_cycles']: - op['trace_events'].append(e) + # Map to other ops + idx = bisect.bisect_right(other_starts, cyc) - 1 + if idx >= 0: + op = other_ops[idx] + if op['start_cycles'] <= cyc <= op['end_cycles']: + op['trace_events'].append(e) return all_ops @@ -563,6 +638,7 @@ def main(): parser.add_argument("--timeline", type=str, nargs='?', const='summary', choices=["summary", "bubbles"], help="Output ASCII art event summary or thread idle bubble analysis (default: summary)") parser.add_argument("--filter", type=str, help="Regex filter matching against the original profile-op line") + parser.add_argument("--device", type=str, help="Device to filter by (e.g. HTP0, HTP0:0) or 'split' to generate separate reports per device") group = parser.add_mutually_exclusive_group() group.add_argument("--head", type=int, help="Limit to first N ops") @@ -586,29 +662,84 @@ def main(): logger.warning(f"Invalid width format '{w}'") final_pmu_name = (args.pmu_name or f"#{args.pmu_index}") if args.pmu_index is not None else None - ops = parse_log(args.logfile, pmu_index=args.pmu_index) + op_filter_re = None if args.filter: try: - filter_re = re.compile(args.filter) + op_filter_re = re.compile(args.filter) except re.error as e: logger.error(f"Invalid regex filter: {e}") sys.exit(1) - ops = [op for op in ops if filter_re.search(op['op_text'])] - - if args.head is not None: - ops = ops[:args.head] - elif args.tail is not None: - ops = ops[-args.tail:] - - if args.timeline: - for op in ops: - if args.timeline == "summary": - print_ascii_summary(op['name'], op['dims'], op['types'], op['usec'], op['cycles'], op['trace_events']) - elif args.timeline == "bubbles": - print_bubbles_timeline(op) + + limit = args.head if args.head is not None else None + device_filter = args.device if (args.device and args.device != "split") else None + ops = parse_log(args.logfile, pmu_index=args.pmu_index, limit=limit, device_filter=device_filter, op_filter_re=op_filter_re) + + if args.device and args.device != "split": + ops = [op for op in ops if device_matches(op['device'], args.device)] + + if args.device == "split": + unique_devices = sorted(list(set(op['device'] for op in ops))) + for dev in unique_devices: + dev_ops = [op for op in ops if device_matches(op['device'], dev)] + + if args.filter: + try: + filter_re = re.compile(args.filter) + except re.error as e: + logger.error(f"Invalid regex filter: {e}") + sys.exit(1) + dev_ops = [op for op in dev_ops if filter_re.search(op['op_text'])] + + if args.head is not None: + dev_ops = dev_ops[:args.head] + elif args.tail is not None: + dev_ops = dev_ops[-args.tail:] + + logger.info("\n=========================================") + logger.info(f" Device: {dev}") + logger.info("=========================================") + + if args.timeline: + for op in dev_ops: + if args.timeline == "summary": + print_ascii_summary(op['name'], op['dims'], op['types'], op['usec'], op['cycles'], op['trace_events']) + elif args.timeline == "bubbles": + print_bubbles_timeline(op) + else: + generate_report(dev_ops, args.top, overrides, args.sort, pmu_name=final_pmu_name) else: - generate_report(ops, args.top, overrides, args.sort, pmu_name=final_pmu_name) + if args.filter: + try: + filter_re = re.compile(args.filter) + except re.error as e: + logger.error(f"Invalid regex filter: {e}") + sys.exit(1) + ops = [op for op in ops if filter_re.search(op['op_text'])] + + if args.head is not None or args.tail is not None: + ops_by_dev = defaultdict(list) + for op in ops: + ops_by_dev[op['device']].append(op) + + filtered_ops = [] + for dev in sorted(ops_by_dev.keys()): + dev_ops = ops_by_dev[dev] + if args.head is not None: + dev_ops = dev_ops[:args.head] + elif args.tail is not None: + dev_ops = dev_ops[-args.tail:] + filtered_ops.extend(dev_ops) + ops = filtered_ops + + if args.timeline: + for op in ops: + if args.timeline == "summary": + print_ascii_summary(op['name'], op['dims'], op['types'], op['usec'], op['cycles'], op['trace_events']) + elif args.timeline == "bubbles": + print_bubbles_timeline(op) + else: + generate_report(ops, args.top, overrides, args.sort, pmu_name=final_pmu_name) if __name__ == "__main__": diff --git a/scripts/snapdragon/ggml-hexagon-trace.py b/scripts/snapdragon/ggml-hexagon-trace.py index 4755adfa133..99bf771b85d 100755 --- a/scripts/snapdragon/ggml-hexagon-trace.py +++ b/scripts/snapdragon/ggml-hexagon-trace.py @@ -20,6 +20,31 @@ r"trace-evt\s+(?P[A-Z_0-9\-]+):\s+thread\s+(?P\d+)\s+info\s+(?P\d+)\s+(?Pstart|stop)\s+(?P\d+)" ) +device_pattern = re.compile(r"\b(HTP\d+(?::\d+)?)\s+(?:profile-op|trace-evt)\b") + + +def extract_device(line): + m = device_pattern.search(line) + if m: + return m.group(1) + return "HTP0" + + +def device_matches(record_device, target_device): + targets = [t.strip() for t in target_device.split(',')] + for target in targets: + if record_device == target: + return True + if record_device.startswith(target + ":"): + return True + return False + + +def get_split_output_path(base_path, device_name): + safe_device = device_name.replace(':', '_') + root, ext = os.path.splitext(base_path) + return f"{root}-{safe_device}{ext}" + def normalize_event_name(evt_type, info=0): if evt_type == "HVX_COMP": @@ -54,7 +79,79 @@ def unwrap(self, raw): return raw + self.high_part -def parse_log(file_path): +class DeviceTimeMapper: + def __init__(self, dev, ops): + self.dev = dev + self.batches = [] + for op in ops: + if op.get('device') == dev and op.get('name') == 'OPBATCH' and op.get('unwrapped_cycles_start') is not None: + cycles = op.get('cycles', 0) + usec = op.get('usec', 0) + start_cyc = op['unwrapped_cycles_start'] + freq = (cycles / usec) if usec > 0 and cycles > 0 else 1000.0 + if freq <= 0: + freq = 1000.0 + self.batches.append({ + 'start_cycles': start_cyc, + 'cycles': cycles, + 'end_cycles': start_cyc + cycles, + 'usec': usec, + 'dur_ns': usec * 1000, + 'freq_mhz': freq, + }) + + self.batches.sort(key=lambda b: b['start_cycles']) + + for i, b in enumerate(self.batches): + if i == 0: + b['start_time_ns'] = 0 + else: + prev = self.batches[i - 1] + idle_cyc = max(0, b['start_cycles'] - prev['end_cycles']) + idle_ns = int(round((idle_cyc / prev['freq_mhz']) * 1000)) + b['start_time_ns'] = prev['start_time_ns'] + prev['dur_ns'] + idle_ns + + self.batch_starts = [b['start_cycles'] for b in self.batches] + + valid_starts = [op['unwrapped_cycles_start'] for op in ops if op.get('device') == dev and op.get('unwrapped_cycles_start') is not None] + self.min_cyc = min(valid_starts) if valid_starts else 0 + if self.batches: + self.default_freq = self.batches[0]['freq_mhz'] + else: + freqs = [op['cycles'] / op['usec'] for op in ops if op.get('device') == dev and op.get('usec', 0) > 0 and op.get('cycles', 0) > 0] + self.default_freq = statistics.mean(freqs) if freqs else 1000.0 + + def get_batch(self, cyc): + if not self.batches: + return None + idx = bisect.bisect_right(self.batch_starts, cyc) - 1 + if idx >= 0: + return self.batches[idx] + return self.batches[0] + + def get_freq(self, cyc=None): + if cyc is not None: + b = self.get_batch(cyc) + if b is not None: + return b['freq_mhz'] + return self.default_freq + + def cycle_to_ns(self, cyc): + if cyc is None: + return 0 + b = self.get_batch(cyc) + if b is not None: + return b['start_time_ns'] + int(round(((cyc - b['start_cycles']) / b['freq_mhz']) * 1000)) + return int(round(((cyc - self.min_cyc) / self.default_freq) * 1000)) + + def dur_cycles_to_ns(self, cyc_start, cyc_dur): + if cyc_dur is None: + return 0 + freq = self.get_freq(cyc_start) + return int(round((cyc_dur / freq) * 1000)) + + +def parse_log(file_path, limit=None, device_filter=None, op_filter_re=None): try: if file_path != "-": f = open(file_path, 'r', encoding='utf-8', errors='ignore') @@ -67,14 +164,25 @@ def parse_log(file_path): all_ops: List[Dict[str, Any]] = [] all_traces: List[Dict[str, Any]] = [] current_op: Optional[Dict[str, Any]] = None - unwrapper = None - trace_unwrapper = None + ops_count_per_device = {} + if device_filter is not None: + for target in device_filter.split(','): + ops_count_per_device[target.strip()] = 0 + limit_reached = False + unwrappers = {} + last_batch_start = {} + trace_unwrappers = {} line_idx = 0 for line in f: line_idx += 1 - if "|" in line and "profile-op" in line: - parts = [p.strip() for p in line.split("|")] + if "profile-op" not in line and "trace-evt" not in line: + continue + device = extract_device(line) + + idx = line.find("profile-op") + if idx != -1 and "|" in line[idx:]: + parts = [p.strip() for p in line[idx:].split("|")] prefix = parts[0] prefix_match = re.search(r"profile-op\s+(?P[A-Z_0-9+]+)", prefix) if not prefix_match: @@ -115,14 +223,18 @@ def parse_log(file_path): if op_name == "OPBATCH": if cycles_start_raw: unwrapped_cycles_start = int(cycles_start_raw) - unwrapper = CycleUnwrapper(unwrapped_cycles_start) - trace_unwrapper = CycleUnwrapper(unwrapped_cycles_start) + unwrappers[device] = CycleUnwrapper(unwrapped_cycles_start) + last_batch_start[device] = unwrapped_cycles_start + for k in list(trace_unwrappers.keys()): + if k[0] == device: + del trace_unwrappers[k] else: - if cycles_start_raw and unwrapper is not None: - unwrapped_cycles_start = unwrapper.unwrap(int(cycles_start_raw)) + if cycles_start_raw: + device_unwrapper = unwrappers.get(device) + if device_unwrapper is not None: + unwrapped_cycles_start = device_unwrapper.unwrap(int(cycles_start_raw)) - idx = line.find("profile-op ") - op_text = line[idx + 11:].strip() if idx != -1 else line.strip() + op_text = re.sub(r"^profile-op\s+", "", line[idx:]).strip() if idx != -1 else line.strip() evt_str = None if types.startswith("evt-cnt "): @@ -142,24 +254,59 @@ def parse_log(file_path): 'cycles_start': int(cycles_start_raw) if cycles_start_raw else None, 'unwrapped_cycles_start': unwrapped_cycles_start, 'trace_events': [], - 'line_num': line_idx + 'line_num': line_idx, + 'device': device } all_ops.append(current_op) + + # Check if matching early exit criteria + matched = False + matched_target = None + if device_filter is not None: + targets = [t.strip() for t in device_filter.split(',')] + for target in targets: + if device == target or device.startswith(target + ":"): + matched = True + matched_target = target + break + else: + matched = True + matched_target = device + + if op_filter_re is not None and not op_filter_re.search(op_text): + matched = False + + if matched: + if matched_target not in ops_count_per_device: + ops_count_per_device[matched_target] = 0 + ops_count_per_device[matched_target] += 1 + + if limit is not None and len(ops_count_per_device) > 0 and all(count >= limit for count in ops_count_per_device.values()): + limit_reached = True + + if limit_reached and op_name == "OPBATCH": + break continue trace_match = trace_pattern.search(line) if trace_match: + thread = int(trace_match.group('thread')) raw_cyc = int(trace_match.group('cycles')) unwrapped_cyc = None - if trace_unwrapper is not None: - unwrapped_cyc = trace_unwrapper.unwrap(raw_cyc) + th_key = (device, thread) + if th_key not in trace_unwrappers: + batch_start = last_batch_start.get(device) + trace_unwrappers[th_key] = CycleUnwrapper(batch_start) + unwrapped_cyc = trace_unwrappers[th_key].unwrap(raw_cyc) all_traces.append({ - 'thread': int(trace_match.group('thread')), + 'thread': thread, 'event': trace_match.group('event'), 'info': int(trace_match.group('info')), 'cycles': raw_cyc, 'unwrapped_cycles': unwrapped_cyc, - 'state': trace_match.group('state') + 'state': trace_match.group('state'), + 'line_num': line_idx, + 'device': device }) f.close() @@ -274,27 +421,24 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path): logger.warning("No operators found after filtering.") return - # Compute average frequency - frequencies = [] - for op in filtered_ops: - if op['usec'] > 0 and op['cycles'] > 0: - frequencies.append(op['cycles'] / op['usec']) - avg_freq_mhz = statistics.mean(frequencies) if frequencies else 1000.0 - if avg_freq_mhz <= 0: - avg_freq_mhz = 1000.0 - # Assign start and end cycles to each operator for op in filtered_ops: op['start_cycles'] = op['unwrapped_cycles_start'] - op['end_cycles'] = op['start_cycles'] + op['cycles'] + op['end_cycles'] = op['start_cycles'] + op['cycles'] if op['start_cycles'] is not None else None - global_min_cyc = min(op['start_cycles'] for op in filtered_ops if op['start_cycles'] is not None) + # Get list of unique devices present in the operations + unique_devices = sorted(list(set(op['device'] for op in filtered_ops))) + device_to_idx = {dev: idx for idx, dev in enumerate(unique_devices)} + time_mappers = {dev: DeviceTimeMapper(dev, filtered_ops) for dev in unique_devices} # Process events completed_events = [] if trace_events: trace_events = sorted(trace_events, key=lambda e: e['unwrapped_cycles']) - one_usec_cycles = max(avg_freq_mhz, 1.0) + + one_usec_cycles = {} + for dev in unique_devices: + one_usec_cycles[dev] = max(time_mappers[dev].get_freq(), 1.0) active_starts = {} for e in trace_events: @@ -303,31 +447,36 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path): info = e['info'] state = e['state'] cyc = e['unwrapped_cycles'] + dev = e['device'] - key = (t, evt, info) + key = (dev, t, evt, info) if state == 'start': # Handle missing stop (start followed by another start) if key in active_starts: - prev_start = active_starts[key] + prev_e = active_starts[key] completed_events.append({ 'thread': t, 'event': evt, 'info': info, - 'start_cyc': prev_start, - 'end_cyc': prev_start + one_usec_cycles, + 'start_cyc': prev_e['unwrapped_cycles'], + 'end_cyc': prev_e['unwrapped_cycles'] + one_usec_cycles.get(dev, 1000.0), + 'line_num': prev_e.get('line_num'), 'missing_stop': True, + 'device': dev }) - active_starts[key] = cyc + active_starts[key] = e elif state == 'stop': if key in active_starts: - start_cyc = active_starts[key] + prev_e = active_starts[key] del active_starts[key] completed_events.append({ 'thread': t, 'event': evt, 'info': info, - 'start_cyc': start_cyc, + 'start_cyc': prev_e['unwrapped_cycles'], 'end_cyc': cyc, + 'line_num': prev_e.get('line_num'), + 'device': dev }) else: # Handle missing start (stop without start) @@ -335,31 +484,36 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path): 'thread': t, 'event': evt, 'info': info, - 'start_cyc': cyc - one_usec_cycles, + 'start_cyc': cyc - one_usec_cycles.get(dev, 1000.0), 'end_cyc': cyc, + 'line_num': e.get('line_num'), 'missing_start': True, + 'device': dev }) # Clear remaining unmatched starts - for key, start_cyc in active_starts.items(): - t, evt, info = key + for key, prev_e in active_starts.items(): + dev, t, evt, info = key completed_events.append({ 'thread': t, 'event': evt, 'info': info, - 'start_cyc': start_cyc, - 'end_cyc': start_cyc + one_usec_cycles, + 'start_cyc': prev_e['unwrapped_cycles'], + 'end_cyc': prev_e['unwrapped_cycles'] + one_usec_cycles.get(dev, 1000.0), + 'line_num': prev_e.get('line_num'), 'missing_stop': True, + 'device': dev }) completed_events.sort(key=lambda e: e['start_cyc']) - # Convert event times to microseconds and apply clamp rounded to 1ns resolution (3 decimals) + # Convert event times to nanoseconds using per-device / per-batch time mapper for e in completed_events: - start_us = (e['start_cyc'] - global_min_cyc) / avg_freq_mhz - dur_us = (e['end_cyc'] - e['start_cyc']) / avg_freq_mhz - e['ts_ns'] = int(round(start_us * 1000)) - e['dur_ns'] = int(round(max(dur_us, 0.1) * 1000)) + dev = e['device'] + tm = time_mappers[dev] + e['ts_ns'] = tm.cycle_to_ns(e['start_cyc']) + dur_ns = tm.dur_cycles_to_ns(e['start_cyc'], e['end_cyc'] - e['start_cyc']) + e['dur_ns'] = max(dur_ns, 100) # Allocate slots (sub-tracks) to prevent overlaps on same virtual track active_slots = defaultdict(list) @@ -368,14 +522,15 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path): evt = e['event'] ts = e['ts_ns'] dur = e['dur_ns'] + dev = e['device'] norm_evt = normalize_event_name(evt, e['info']) if norm_evt == "DMA": - track_key = (t, "DMA") + track_key = (dev, t, "DMA") elif t == 10: - track_key = (t, "HMX") + track_key = (dev, t, "HMX") else: - track_key = (t, "HVX") + track_key = (dev, t, "HVX") slots = active_slots[track_key] allocated_slot = -1 @@ -395,6 +550,7 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path): t = e['thread'] evt = e['event'] slot = e['slot'] + dev = e['device'] norm_evt = normalize_event_name(evt, e['info']) if norm_evt == "DMA": @@ -408,56 +564,69 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path): evt_id = 2 t_sort = 1 if t == 10 else t + 2 + dev_idx = device_to_idx[dev] + # Unique UUID for each sub-track if t == 10: - uuid = 20 # HMX thread track UUID + uuid = dev_idx * 10000000 + 20 # HMX thread track UUID else: - uuid = int(t_sort * 1000000 + evt_id * 1000 + slot) + uuid = int(dev_idx * 10000000 + t_sort * 1000000 + evt_id * 1000 + slot) e['uuid'] = uuid - used_tracks[uuid] = (t, track_evt, slot) + used_tracks[uuid] = (dev, t, track_evt, slot) with open(output_path, "wb") as f: - # Define Process with EXPLICIT child sorting - proc_desc = make_process_descriptor(1, "HTP NPU") - proc_packet = make_trace_packet(0, track_descriptor=make_track_descriptor(1, process=proc_desc, child_ordering=3)) - write_trace_packet_to_file(f, proc_packet) - - # Define Operators Track (UUID = 2) as a thread track at rank 1, tid 8 - op_thread_desc = make_thread_descriptor(1, 8, "Ops", sort_index=1) - op_packet = make_trace_packet(0, track_descriptor=make_track_descriptor(2, parent_uuid=1, thread=op_thread_desc)) - write_trace_packet_to_file(f, op_packet) - - # Define HMX Thread Track (UUID = 20) at rank 2, tid 9 - hmx_thread_desc = make_thread_descriptor(1, 9, "HMX", sort_index=2) - hmx_packet = make_trace_packet(0, track_descriptor=make_track_descriptor(20, parent_uuid=1, thread=hmx_thread_desc)) - write_trace_packet_to_file(f, hmx_packet) - - # Define Thread Tracks (T0, T1, ..., T9) - unique_threads = sorted(list(set(t for (t, _, _) in used_tracks.values() if t != 10))) - for t in unique_threads: - thread_uuid = 10 + t - thread_name = f"T{t}" - # Sort order starts from index 3 (T0 -> 3, T1 -> 4, etc.) - sort_index = 3 + t - tid = 10 + t - thread_desc = make_thread_descriptor(1, tid, thread_name, sort_index=sort_index) - thread_packet = make_trace_packet(0, track_descriptor=make_track_descriptor( - thread_uuid, - parent_uuid=1, - thread=thread_desc, - sibling_order_rank=sort_index, - child_ordering=3 # Explicit child sorting for sub-tracks - )) - write_trace_packet_to_file(f, thread_packet) + for dev in unique_devices: + dev_idx = device_to_idx[dev] + pid = dev_idx + 1 + proc_uuid = dev_idx * 10000000 + 1 + + # Define Process with EXPLICIT child sorting + proc_name = dev + proc_desc = make_process_descriptor(pid, proc_name) + proc_packet = make_trace_packet(0, track_descriptor=make_track_descriptor(proc_uuid, process=proc_desc, child_ordering=3)) + write_trace_packet_to_file(f, proc_packet) + + # Define Operators Track as a thread track + op_track_uuid = dev_idx * 10000000 + 2 + op_tid = pid * 100 + 8 + op_thread_desc = make_thread_descriptor(pid, op_tid, "Ops", sort_index=1) + op_packet = make_trace_packet(0, track_descriptor=make_track_descriptor(op_track_uuid, parent_uuid=proc_uuid, thread=op_thread_desc)) + write_trace_packet_to_file(f, op_packet) + + # Define HMX Thread Track at rank 2 + hmx_track_uuid = dev_idx * 10000000 + 20 + hmx_tid = pid * 100 + 9 + hmx_thread_desc = make_thread_descriptor(pid, hmx_tid, "HMX", sort_index=2) + hmx_packet = make_trace_packet(0, track_descriptor=make_track_descriptor(hmx_track_uuid, parent_uuid=proc_uuid, thread=hmx_thread_desc)) + write_trace_packet_to_file(f, hmx_packet) + + # Define Thread Tracks (T0, T1, ..., T9) for this device + dev_used_tracks = {uuid: val for uuid, val in used_tracks.items() if val[0] == dev} + unique_threads = sorted(list(set(t for (_, t, _, _) in dev_used_tracks.values() if t != 10))) + for t in unique_threads: + thread_uuid = dev_idx * 10000000 + 10 + t + thread_name = f"T{t}" + sort_index = 3 + t + tid = pid * 100 + 10 + t + thread_desc = make_thread_descriptor(pid, tid, thread_name, sort_index=sort_index) + thread_packet = make_trace_packet(0, track_descriptor=make_track_descriptor( + thread_uuid, + parent_uuid=proc_uuid, + thread=thread_desc, + sibling_order_rank=sort_index, + child_ordering=3 # Explicit child sorting for sub-tracks + )) + write_trace_packet_to_file(f, thread_packet) # Define Track descriptors for sub-tracks parented to thread tracks for uuid in sorted(used_tracks.keys()): - if uuid == 20: + dev, t, evt, slot = used_tracks[uuid] + dev_idx = device_to_idx[dev] + if t == 10: continue - t, evt, slot = used_tracks[uuid] name = f"T{t} {evt}" rank = 0 if evt == "HVX" else 1 - parent_thread_uuid = 10 + t + parent_thread_uuid = dev_idx * 10000000 + 10 + t # Sibling merge behavior: 1 (SIBLING_MERGE_BEHAVIOR_BY_TRACK_NAME) track_desc = make_track_descriptor( uuid=uuid, @@ -470,15 +639,18 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path): write_trace_packet_to_file(f, track_packet) # Emit Operators - last_op_end_ns = 0 + last_op_end_ns = defaultdict(int) for op in filtered_ops: - op_start_ns = int(round(((op['start_cycles'] - global_min_cyc) / avg_freq_mhz) * 1000)) - op_dur_ns = int(round((op['cycles'] / avg_freq_mhz) * 1000)) + dev = op['device'] + dev_idx = device_to_idx[dev] + tm = time_mappers[dev] + op_start_ns = tm.cycle_to_ns(op['start_cycles']) + op_dur_ns = tm.dur_cycles_to_ns(op['start_cycles'], op['cycles']) if op['name'] != "OPBATCH": - if op_start_ns < last_op_end_ns: - op_start_ns = last_op_end_ns + if op_start_ns < last_op_end_ns[dev]: + op_start_ns = last_op_end_ns[dev] clamped_dur = max(op_dur_ns, 100) # Clamp to 100ns (0.1us) - last_op_end_ns = op_start_ns + clamped_dur + last_op_end_ns[dev] = op_start_ns + clamped_dur else: clamped_dur = max(op_dur_ns, 100) @@ -495,24 +667,41 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path): if 'evt' in op and op['evt']: debug_annots.append(make_debug_annotation("evt", string_val=op['evt'])) + op_track_uuid = dev_idx * 10000000 + 2 + # Slice Begin - evt_begin = make_track_event(1, 2, name=f"{op['name']} ({op['dims']})", category="operator", debug_annotations=debug_annots) + evt_begin = make_track_event(1, op_track_uuid, name=f"{op['name']} ({op['dims']})", category="operator", debug_annotations=debug_annots) packet_begin = make_trace_packet(op_start_ns, track_event=evt_begin) write_trace_packet_to_file(f, packet_begin) # Slice End - evt_end = make_track_event(2, 2) + evt_end = make_track_event(2, op_track_uuid) packet_end = make_trace_packet(op_start_ns + clamped_dur, track_event=evt_end) write_trace_packet_to_file(f, packet_end) # Emit Thread Trace Events for e in completed_events: norm_name = normalize_event_name(e['event'], e['info']) - name = f"DMA {e['info']}" if norm_name == "DMA" else norm_name + if norm_name == "DMA": + name = f"DMA {e['info']}" + elif norm_name == "FENCE": + name = f"FENCE {e['info']}" if e.get('info') is not None and e['info'] != 0 else "FENCE" + else: + name = norm_name + if e.get('missing_start') or e.get('missing_stop'): name += "!" debug_annots = [] + if 'line_num' in e and e['line_num'] is not None: + debug_annots.append(make_debug_annotation("line", int_val=e['line_num'])) + if norm_name == "FENCE" and e.get('info') is not None: + debug_annots.append(make_debug_annotation("seq", int_val=e['info'])) + elif norm_name == "DMA" and e.get('info') is not None: + debug_annots.append(make_debug_annotation("channel", int_val=e['info'])) + elif e.get('info') is not None and e['info'] != 0: + debug_annots.append(make_debug_annotation("info", int_val=e['info'])) + if e.get('missing_start'): debug_annots.append(make_debug_annotation("missing_start", string_val="true")) if e.get('missing_stop'): @@ -536,6 +725,7 @@ def main(): parser.add_argument("logfile", help="Path to hex-log profile file") parser.add_argument("-o", "--output", default="optrace.perfetto-trace", help="Output trace file path (default: optrace.perfetto-trace)") parser.add_argument("--filter", type=str, help="Regex filter matching against the original profile-op line") + parser.add_argument("--device", type=str, help="Device to filter by (e.g. HTP0, HTP0:0) or 'split' to generate separate files per device") group = parser.add_mutually_exclusive_group() group.add_argument("--head", type=int, help="Limit to first N ops") @@ -544,7 +734,21 @@ def main(): args = parser.parse_args() logging.basicConfig(level=logging.INFO, format='%(message)s') - ops, traces = parse_log(args.logfile) + op_filter_re = None + if args.filter: + try: + op_filter_re = re.compile(args.filter) + except re.error as e: + logger.error(f"Invalid regex filter: {e}") + sys.exit(1) + + limit = args.head if args.head is not None else None + device_filter = args.device if (args.device and args.device != "split") else None + ops, traces = parse_log(args.logfile, limit=limit, device_filter=device_filter, op_filter_re=op_filter_re) + + if args.device and args.device != "split": + ops = [op for op in ops if device_matches(op['device'], args.device)] + traces = [t for t in traces if device_matches(t['device'], args.device)] if args.filter: try: @@ -554,35 +758,60 @@ def main(): sys.exit(1) ops = [op for op in ops if filter_re.search(op['op_text'])] - if args.head is not None: - ops = ops[:args.head] - elif args.tail is not None: - ops = ops[-args.tail:] + if args.head is not None or args.tail is not None: + ops_by_dev = defaultdict(list) + for op in ops: + ops_by_dev[op['device']].append(op) + + filtered_ops = [] + for dev in sorted(ops_by_dev.keys()): + dev_ops = ops_by_dev[dev] + if args.head is not None: + dev_ops = dev_ops[:args.head] + elif args.tail is not None: + dev_ops = dev_ops[-args.tail:] + filtered_ops.extend(dev_ops) + ops = filtered_ops if args.filter or args.head is not None or args.tail is not None: - valid_ranges = [] + # Group valid ranges by device + valid_ranges_by_dev = defaultdict(list) for op in ops: start_cyc = op['unwrapped_cycles_start'] end_cyc = start_cyc + op['cycles'] if start_cyc is not None else None if start_cyc is not None and end_cyc is not None: - valid_ranges.append((start_cyc, end_cyc)) + valid_ranges_by_dev[op['device']].append((start_cyc, end_cyc)) - valid_ranges.sort(key=lambda r: r[0]) - range_starts = [r[0] for r in valid_ranges] + for dev in valid_ranges_by_dev: + valid_ranges_by_dev[dev].sort(key=lambda r: r[0]) + + range_starts_by_dev = {dev: [r[0] for r in ranges] for dev, ranges in valid_ranges_by_dev.items()} filtered_traces = [] for e in traces: cyc = e['unwrapped_cycles'] if cyc is None: continue + dev = e['device'] + range_starts = range_starts_by_dev.get(dev) + if not range_starts: + continue idx = bisect.bisect_right(range_starts, cyc) - 1 if idx >= 0: - start, end = valid_ranges[idx] + start, end = valid_ranges_by_dev[dev][idx] if start <= cyc <= end: filtered_traces.append(e) traces = filtered_traces - generate_perfetto_trace(ops, traces, args.output) + if args.device == "split": + unique_devices = sorted(list(set(op['device'] for op in ops))) + for dev in unique_devices: + dev_ops = [op for op in ops if device_matches(op['device'], dev)] + dev_traces = [t for t in traces if device_matches(t['device'], dev)] + out_path = get_split_output_path(args.output, dev) + generate_perfetto_trace(dev_ops, dev_traces, out_path) + else: + generate_perfetto_trace(ops, traces, args.output) if __name__ == "__main__": diff --git a/scripts/snapdragon/run.py b/scripts/snapdragon/run.py new file mode 100755 index 00000000000..d639b16b825 --- /dev/null +++ b/scripts/snapdragon/run.py @@ -0,0 +1,405 @@ +#!/usr/bin/env python3 +# +# Run llama.cpp tools on Snapdragon devices (natively, via ADB, or SSH). +# + +import sys +import os +import argparse +import subprocess +import platform +import shlex +import logging + +logger = logging.getLogger("run") + + +def parse_target(target_str): + if not target_str: + return None, None + if target_str.startswith("adb") or target_str.startswith("android"): + parts = target_str.split(":", 1) + serial = parts[1] if len(parts) > 1 else None + return "android", serial + elif target_str.startswith("lnx") or target_str.startswith("linux") or target_str.startswith("ubuntu"): + parts = target_str.split(":", 1) + host = parts[1] if len(parts) > 1 else None + return "linux", host + elif target_str in ("wos", "windows"): + return "windows", None + else: + return None, None + + +def shlex_join(args_list): + if hasattr(shlex, 'join'): + return shlex.join(args_list) + import pipes + return " ".join(pipes.quote(x) for x in args_list) + + +def main(): + logging.basicConfig(level=logging.INFO, format='%(message)s') + # Split arguments at '--' + if '--' in sys.argv: + idx = sys.argv.index('--') + run_args = sys.argv[1:idx] + cmd_args = sys.argv[idx + 1:] + else: + run_args = sys.argv[1:] + cmd_args = [] + + parser = argparse.ArgumentParser( + description="Unified runner for llama.cpp tools on Snapdragon (natively, via ADB, or via SSH)." + ) + parser.add_argument("--target", help="Execution target (e.g. android[:serial]/adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, windows/wos) (default: local run)") + parser.add_argument("--target-dir", help="Target directory on the device (default: /data/local/tmp/llama.cpp for Android, ~/llama.cpp for Linux)") + parser.add_argument("--install-dir", help="Install directory name (defaults to pkg-TARGET or pkg-TARGET-dbg prefix based on target)") + parser.add_argument("--debug", action="store_true", help="Use debug build (defaults to pkg-TARGET-dbg folder)") + parser.add_argument("--devices", "--device", "-d", help="Select execution devices (split into NPU and OpenCL GPUs automatically, default: HTP0)") + parser.add_argument("--verbose", help="Verbose level (enables both Hexagon and OpenCL kernel cache debugging)") + parser.add_argument("--profile", help="Profiling flag (enables Hexagon profiling and OpenCL autotuning)") + parser.add_argument("--sched-debug", action="store_true", help="Enable GGML/llama.cpp scheduler debug output (GGML_SCHED_DEBUG=2)") + parser.add_argument("--mtmd-device", help="Specify the backend device ID for Multi-Threaded Multi-Device setup (MTMD_BACKEND_DEVICE)") + + # Hexagon specific parameters + parser.add_argument("--hex-verbose", help="Enable verbose logging (GGML_HEXAGON_VERBOSE)") + parser.add_argument("--hex-profile", help="Enable NPU/Hexagon profiling and performance metrics print (GGML_HEXAGON_PROFILE)") + parser.add_argument("--hex-nhvx", help="Number of HVX units to use (GGML_HEXAGON_NHVX)") + parser.add_argument("--hex-nhmx", help="Number of HMX units to use. 0 disables HMX power-up (GGML_HEXAGON_NHMX)") + parser.add_argument("--hex-hostbuf", help="Enable host buffers (GGML_HEXAGON_HOSTBUF)") + parser.add_argument("--hex-opbatch", help="Maximum number of operations to batch into a single HTP execution (GGML_HEXAGON_OPBATCH)") + parser.add_argument("--hex-opqueue", help="Size of the asynchronous NPU operation queue (GGML_HEXAGON_OPQUEUE)") + parser.add_argument("--hex-oppoll", default="1", help="Enable (1) or Disable (0) polling for NPU opbatch completion (GGML_HEXAGON_OPPOLL) (default: 1)") + parser.add_argument("--hex-opfilter", help="Regex pattern to filter/select which operators are offloaded to NPU (GGML_HEXAGON_OPFILTER)") + parser.add_argument("--hex-opfusion", help="NPU graph node fusion optimization level (0: disabled, 1: enabled) (GGML_HEXAGON_OPFUSION)") + parser.add_argument("--hex-vmem", help="Maximum NPU VMEM size limit in MB to allocate (GGML_HEXAGON_VMEM)") + parser.add_argument("--hex-mbuf", help="Maximum host buffer size limit in MB to allocate (GGML_HEXAGON_MBUF)") + parser.add_argument("--hex-mm-select", help="Select MUL_MAT and MUL_MAT_ID kernel (GGML_HEXAGON_MM_SELECT) 3:HMX,2:HVX-tiled,1:HVX-flat,0:disable") + parser.add_argument("--hex-fa-select", help="Select Flash Attention kernel (GGML_HEXAGON_FA_SELECT) 2:HMX,1:HVX,0:disable") + parser.add_argument("--hex-ar-select", help="Select All-Reduce kernel (GGML_HEXAGON_AR_SELECT) 1:enable,0:disable") + parser.add_argument("--hex-etm", help="Enable Embedded Trace Macrocell hardware tracing / trace logging (GGML_HEXAGON_ETM)") + parser.add_argument("--hex-arch", help="Target Hexagon NPU architecture version override (v73, v75, v79, v81, etc.) (GGML_HEXAGON_ARCH)") + parser.add_argument("--hex-optrace", help="Trace buffer size in number of records (GGML_HEXAGON_OPTRACE)") + + # OpenCL specific parameters + parser.add_argument("--cl-platform", help="Select OpenCL platform name/regex (e.g. Qualified Qualcomm OpenCL platform) (GGML_OPENCL_PLATFORM)") + parser.add_argument("--cl-device", help="Select OpenCL device name/regex (e.g. Adreno GPU) (GGML_OPENCL_DEVICE)") + parser.add_argument("--cl-opfilter", help="Regex pattern to filter/select which operators are offloaded to OpenCL (GGML_OPENCL_OPFILTER)") + parser.add_argument("--cl-disable-fusion", action="store_true", help="Disable OpenCL kernel fusion optimizations (GGML_OPENCL_DISABLE_FUSION)") + parser.add_argument("--cl-cache-dir", help="Directory path to store compiled OpenCL program binaries (GGML_OPENCL_KERNEL_CACHE_DIR)") + parser.add_argument("--cl-cache-debug", help="Enable verbose debugging logs for the kernel caching system (GGML_OPENCL_KERNEL_CACHE_DEBUG)") + parser.add_argument("--cl-fa-tune", action="store_true", help="Enable automatic Flash Attention kernel autotuning (GGML_OPENCL_FA_TUNE)") + parser.add_argument("--cl-adreno-xmem", action="store_true", help="Enforce matmul using texture/image (xmem) memory paths on Adreno GPUs (GGML_OPENCL_ADRENO_XMEM_GEMM)") + parser.add_argument("--cl-adreno-large-buffer", action="store_true", help="Allow allocating larger buffer sizes on Adreno GPUs (GGML_OPENCL_ADRENO_USE_LARGE_BUFFER)") + + args = parser.parse_args(run_args) + + if not cmd_args: + parser.print_help() + logger.error("\nError: No command specified after '--'") + sys.exit(1) + + target_type = None + target_val = None + target_prefix = None + if args.target: + target_type, target_val = parse_target(args.target) + if not target_type: + logger.error(f"Error: Invalid target format '{args.target}'. Must be android[:serial]/adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, or windows/wos.") + sys.exit(1) + target_prefix = args.target.split(":", 1)[0] + + # Resolve install directory + install_dir = args.install_dir + if not install_dir: + if target_prefix: + suffix = "-dbg" if args.debug else "" + install_dir = f"pkg-{target_prefix}{suffix}" + else: + # Smart branch folder detection for local run if default is not set + prefixes = ("wos", "windows", "lnx", "linux", "ubuntu", "adb", "android") + suffixes = ("-dbg", "") if args.debug else ("", "-dbg") + found = False + for suffix in suffixes: + for prefix in prefixes: + test_path = f"./pkg-{prefix}{suffix}/llama.cpp" + if os.path.exists(test_path): + install_dir = f"pkg-{prefix}{suffix}" + found = True + break + if found: + break + if not install_dir: + install_dir = "pkg-android" # Fallback default + + # Host side package path + package_path = os.path.join(install_dir, "llama.cpp") + + # Environment variables to map + env_vars = {} + + def set_env(env_name, opt_val): + if opt_val is not None: + env_vars[env_name] = str(opt_val) + elif env_name in os.environ: + env_vars[env_name] = os.environ[env_name] + + # Resolve and filter devices (HTP vs OpenCL) + devices_val = args.devices if args.devices is not None else "HTP0" + if devices_val.isdigit(): + hex_devices = devices_val + cl_device = "" + else: + parts = [p.strip() for p in devices_val.split(",")] + # Any device containing "htp" is Hexagon, rest is OpenCL + hex_parts = [p for p in parts if "htp" in p.lower()] + cl_parts = [p for p in parts if "htp" not in p.lower()] + hex_devices = ",".join(hex_parts) + cl_device = ",".join(cl_parts) + + # Set Hexagon devices + if hex_devices: + env_vars["GGML_HEXAGON_DEVICES"] = hex_devices + elif "GGML_HEXAGON_DEVICES" in os.environ: + env_vars["GGML_HEXAGON_DEVICES"] = os.environ["GGML_HEXAGON_DEVICES"] + + # Set OpenCL device (unless overridden by --cl-device) + final_cl_device = args.cl_device if args.cl_device is not None else cl_device + if final_cl_device: + env_vars["GGML_OPENCL_DEVICE"] = final_cl_device + elif "GGML_OPENCL_DEVICE" in os.environ: + env_vars["GGML_OPENCL_DEVICE"] = os.environ["GGML_OPENCL_DEVICE"] + + # Map shared & backend-specific parameters with correct overrides + + # Verbose logging mapping + hex_verbose_val = args.hex_verbose if args.hex_verbose is not None else args.verbose + set_env("GGML_HEXAGON_VERBOSE", hex_verbose_val) + + cl_cache_debug_val = args.cl_cache_debug if args.cl_cache_debug is not None else args.verbose + set_env("GGML_OPENCL_KERNEL_CACHE_DEBUG", cl_cache_debug_val) + + # Profiling mapping + hex_profile_val = args.hex_profile if args.hex_profile is not None else args.profile + set_env("GGML_HEXAGON_PROFILE", hex_profile_val) + + if args.cl_fa_tune or args.profile is not None: + env_vars["GGML_OPENCL_FA_TUNE"] = "1" + elif "GGML_OPENCL_FA_TUNE" in os.environ: + env_vars["GGML_OPENCL_FA_TUNE"] = os.environ["GGML_OPENCL_FA_TUNE"] + + # Other Hexagon environment variables + set_env("GGML_HEXAGON_NHVX", args.hex_nhvx) + set_env("GGML_HEXAGON_NHMX", args.hex_nhmx) + set_env("GGML_HEXAGON_HOSTBUF", args.hex_hostbuf) + set_env("GGML_HEXAGON_OPBATCH", args.hex_opbatch) + set_env("GGML_HEXAGON_OPQUEUE", args.hex_opqueue) + set_env("GGML_HEXAGON_OPPOLL", args.hex_oppoll) + set_env("GGML_HEXAGON_OPFILTER", args.hex_opfilter) + set_env("GGML_HEXAGON_OPFUSION", args.hex_opfusion) + set_env("GGML_HEXAGON_VMEM", args.hex_vmem) + set_env("GGML_HEXAGON_MBUF", args.hex_mbuf) + set_env("GGML_HEXAGON_MM_SELECT", args.hex_mm_select) + set_env("GGML_HEXAGON_FA_SELECT", args.hex_fa_select) + set_env("GGML_HEXAGON_AR_SELECT", args.hex_ar_select) + set_env("GGML_HEXAGON_ETM", args.hex_etm) + set_env("GGML_HEXAGON_ARCH", args.hex_arch) + set_env("GGML_HEXAGON_OPTRACE", args.hex_optrace) + set_env("MTMD_BACKEND_DEVICE", args.mtmd_device) + + # OpenCL environment variables + set_env("GGML_OPENCL_PLATFORM", args.cl_platform) + set_env("GGML_OPENCL_OPFILTER", args.cl_opfilter) + set_env("GGML_OPENCL_KERNEL_CACHE_DIR", args.cl_cache_dir) + + if args.cl_disable_fusion: + env_vars["GGML_OPENCL_DISABLE_FUSION"] = "1" + elif "GGML_OPENCL_DISABLE_FUSION" in os.environ: + env_vars["GGML_OPENCL_DISABLE_FUSION"] = os.environ["GGML_OPENCL_DISABLE_FUSION"] + + if args.cl_adreno_xmem: + env_vars["GGML_OPENCL_ADRENO_XMEM_GEMM"] = "1" + elif "GGML_OPENCL_ADRENO_XMEM_GEMM" in os.environ: + env_vars["GGML_OPENCL_ADRENO_XMEM_GEMM"] = os.environ["GGML_OPENCL_ADRENO_XMEM_GEMM"] + + if args.cl_adreno_large_buffer: + env_vars["GGML_OPENCL_ADRENO_USE_LARGE_BUFFER"] = "1" + elif "GGML_OPENCL_ADRENO_USE_LARGE_BUFFER" in os.environ: + env_vars["GGML_OPENCL_ADRENO_USE_LARGE_BUFFER"] = os.environ["GGML_OPENCL_ADRENO_USE_LARGE_BUFFER"] + + if args.sched_debug: + env_vars["GGML_SCHED_DEBUG"] = "2" + + # Resolve executable path + executable = cmd_args[0] + known_binaries = ["llama-cli", "llama-bench", "llama-completion", "llama-mtmd-cli", "test-backend-ops"] + if executable in known_binaries: + if target_type in ("android", "linux"): + resolved_exec = f"./bin/{executable}" + else: + if platform.system() == "Windows": + resolved_exec = os.path.normpath(os.path.join(package_path, "bin", f"{executable}.exe")) + else: + resolved_exec = os.path.normpath(os.path.join(package_path, "bin", executable)) + cmd_args[0] = resolved_exec + + # Infer device string to pass to the tool + basename = os.path.basename(executable) + if basename.endswith(".exe"): + basename = basename[:-4] + + device_val = None + if basename == "test-backend-ops": + for i in range(len(cmd_args)): + if cmd_args[i] in ("-p", "--params") and i + 1 < len(cmd_args): + val = cmd_args[i + 1] + new_val = "" + for j, char in enumerate(val): + if char in ('[', ']'): + if j > 0 and val[j - 1] == '\\': + new_val += char + else: + new_val += '\\' + char + else: + new_val += char + cmd_args[i + 1] = new_val + + has_b = any(arg == "-b" for arg in cmd_args) + if not has_b: + if args.devices: + if args.devices.isdigit(): + n = int(args.devices) + device_val = ",".join(f"HTP{i}" for i in range(n)) + else: + device_val = args.devices + elif "D" in os.environ: + device_val = os.environ["D"] + elif "DEVICE" in os.environ: + device_val = os.environ["DEVICE"] + else: + device_val = "HTP0" + if device_val: + cmd_args += ["-b", device_val] + else: + has_device = any(arg.startswith("--device") for arg in cmd_args) + if not has_device: + if args.devices: + if args.devices.isdigit(): + n = int(args.devices) + device_val = ",".join(f"HTP{i}" for i in range(n)) + else: + device_val = args.devices + elif "D" in os.environ: + device_val = os.environ["D"] + elif "DEVICE" in os.environ: + device_val = os.environ["DEVICE"] + else: + device_val = "HTP0" + if device_val: + cmd_args += ["--device", device_val] + + # Automatically add -v to known llama tools if sched-debug, verbose, or profile are set + verbose_trigger = ( + args.sched_debug + or args.verbose is not None + or args.profile is not None + or args.hex_verbose is not None + or args.hex_profile is not None + or args.hex_optrace is not None + ) + if verbose_trigger and basename in ("llama-cli", "llama-completion", "llama-bench", "llama-server", "llama-mtmd-cli"): + if "-v" not in cmd_args and "--verbose" not in cmd_args: + cmd_args.append("-v") + + # Inject defaults for llama-cli, llama-completion, and llama-server if not overridden by the user + if basename in ("llama-cli", "llama-completion", "llama-server"): + if "-ngl" not in cmd_args and "--n-gpu-layers" not in cmd_args: + cmd_args += ["-ngl", "99"] + if "--ubatch-size" not in cmd_args and "-ub" not in cmd_args: + cmd_args += ["--ubatch-size", "1024"] + if "-fa" not in cmd_args and "--flash-attn" not in cmd_args: + cmd_args += ["-fa", "on"] + + if basename in ("llama-cli", "llama-completion", "llama-server", "llama-bench"): + if "-t" not in cmd_args and "--threads" not in cmd_args: + cmd_args += ["-t", "6"] + + # Resolve target directory on device + target_dir = args.target_dir + if not target_dir: + target_dir = "/data/local/tmp/llama.cpp" if target_type == "android" else "~/llama.cpp" + + if target_type == "android": + # Run via ADB + adb_base = ["adb"] + if target_val: # serial + adb_base += ["-s", target_val] + + env_parts = [ + "LD_LIBRARY_PATH=./lib", + "ADSP_LIBRARY_PATH=./lib" + ] + for k, v in env_vars.items(): + env_parts.append(f"{k}={v}") + env_str = " ".join(env_parts) + + cmd_str = shlex_join(cmd_args) + adb_shell_cmd = f"cd {target_dir} && ulimit -c unlimited && {env_str} {cmd_str}" + full_cmd = adb_base + ["shell", adb_shell_cmd] + + logger.info(f"+ {' '.join(full_cmd)}") + res = subprocess.run(full_cmd) + sys.exit(res.returncode) + + elif target_type == "linux": + ssh_host = target_val + if not ssh_host: + logger.error("Error: SSH host not specified in target (e.g. use linux:user@host, lnx:user@host, or ubuntu:user@host). Cannot execute.") + sys.exit(1) + + # Linux remote run via SSH + env_parts = [ + "LD_LIBRARY_PATH=./lib", + "ADSP_LIBRARY_PATH=./lib" + ] + for k, v in env_vars.items(): + env_parts.append(f"{k}={v}") + env_str = " ".join(env_parts) + + cmd_str = shlex_join(cmd_args) + ssh_shell_cmd = f"cd {target_dir} && ulimit -c unlimited && {env_str} {cmd_str}" + full_cmd = ["ssh", ssh_host, ssh_shell_cmd] + + logger.info(f"+ {' '.join(full_cmd)}") + res = subprocess.run(full_cmd) + sys.exit(res.returncode) + + elif target_type == "windows": + logger.info("Windows target execution is currently a stub.") + sys.exit(0) + + else: + # Run locally + local_env = os.environ.copy() + lib_dir = os.path.normpath(os.path.join(package_path, "lib")) + local_env["ADSP_LIBRARY_PATH"] = lib_dir + if platform.system() == "Windows": + local_env["PATH"] = lib_dir + os.path.pathsep + local_env.get("PATH", "") + else: + local_env["LD_LIBRARY_PATH"] = lib_dir + os.path.pathsep + local_env.get("LD_LIBRARY_PATH", "") + + for k, v in env_vars.items(): + local_env[k] = v + + logger.info(f"+ {shlex_join(cmd_args)}") + res = subprocess.run(cmd_args, env=local_env) + sys.exit(res.returncode) + + +if __name__ == "__main__": + try: + main() + except KeyboardInterrupt: + logger.info("\nInterrupted by user.") + sys.exit(130) diff --git a/scripts/snapdragon/sdk.py b/scripts/snapdragon/sdk.py new file mode 100644 index 00000000000..bb3cb77b2ad --- /dev/null +++ b/scripts/snapdragon/sdk.py @@ -0,0 +1,62 @@ +import os +from pathlib import Path + + +SDK_CONFIGS = ( + { + "name": "Hexagon SDK", + "repo": "snapdragon-toolchain/hexagon-sdk", + "default_version": "6.6.0.0", + "parent_dir": "Hexagon_SDK", + "archive_prefix": "hexagon-sdk-v", + "markers": ("hexagon_sdk.json",), + }, + { + "name": "OpenCL SDK", + "repo": "snapdragon-toolchain/opencl-sdk", + "default_version": "2.3.2", + "parent_dir": "OpenCL_SDK", + "archive_prefix": "adreno-opencl-sdk-v", + "markers": ("include/CL", "lib/OpenCL.lib"), + }, +) + + +def is_valid_sdk(config, target_dir): + return target_dir.is_dir() and all((target_dir / marker).exists() for marker in config["markers"]) + + +def get_hexagon_tools_dir(hexagon_dir): + tools_parent = hexagon_dir / "tools" / "HEXAGON_Tools" + if not tools_parent.is_dir(): + raise RuntimeError(f"Expected Hexagon tools directory in {tools_parent}") + tools_dirs = [path for path in tools_parent.iterdir() if path.is_dir()] + if len(tools_dirs) != 1: + raise RuntimeError(f"Expected one Hexagon tools directory in {tools_parent}") + return tools_dirs[0] + + +def validate_windows_sdks(): + hexagon_config, opencl_config = SDK_CONFIGS + hexagon_dir = os.environ.get("HEXAGON_SDK_ROOT") + tools_dir = os.environ.get("HEXAGON_TOOLS_ROOT") + opencl_dir = os.environ.get("OPENCL_SDK_ROOT") + missing = [] + + expected_tools_dir = None + if not hexagon_dir or not is_valid_sdk(hexagon_config, Path(hexagon_dir)): + missing.append("HEXAGON_SDK_ROOT") + else: + try: + expected_tools_dir = get_hexagon_tools_dir(Path(hexagon_dir)) + except RuntimeError: + pass + if not tools_dir or not expected_tools_dir or Path(tools_dir) != expected_tools_dir: + missing.append("HEXAGON_TOOLS_ROOT") + if not opencl_dir or not is_valid_sdk(opencl_config, Path(opencl_dir)): + missing.append("OPENCL_SDK_ROOT") + if missing: + raise RuntimeError( + f"Missing or invalid Windows SDK paths: {', '.join(missing)}. " + "Run scripts/snapdragon/setup-sdk.py first." + ) diff --git a/scripts/snapdragon/setup-sdk.py b/scripts/snapdragon/setup-sdk.py new file mode 100644 index 00000000000..ad828c079a6 --- /dev/null +++ b/scripts/snapdragon/setup-sdk.py @@ -0,0 +1,233 @@ +#!/usr/bin/env python3 +# +# Install Windows on Snapdragon SDKs for llama.cpp. +# + +import sys +import os +import argparse +import shutil +import logging +import json +import hashlib +import tarfile +import tempfile +from pathlib import Path +from urllib.error import HTTPError, URLError +from urllib.request import Request, urlopen + +from sdk import SDK_CONFIGS, get_hexagon_tools_dir, is_valid_sdk + + +logger = logging.getLogger("setup_sdk") + +DEFAULT_SDK_BASE_DIR = r"C:\Qualcomm" + + +def get_sdk_releases(config): + request = Request( + f"https://api.github.com/repos/{config['repo']}/releases?per_page=100", + headers={"Accept": "application/vnd.github+json", "User-Agent": "llama.cpp"}, + ) + try: + with urlopen(request, timeout=30) as response: + releases = json.load(response) + except (HTTPError, URLError, TimeoutError) as err: + raise RuntimeError(f"Cannot query {config['name']} releases: {err}") from err + + result = [] + for release in releases: + if release["draft"] or release["prerelease"]: + continue + version = release["tag_name"].removeprefix("v") + archive_name = f"{config['archive_prefix']}{version}-arm64-wos.tar.xz" + for asset in release["assets"]: + if asset["name"] != archive_name: + continue + result.append({ + "version": version, + "name": asset["name"], + "url": asset["browser_download_url"], + "sha256": (asset.get("digest") or "").removeprefix("sha256:"), + }) + return result + + +def list_sdk_releases(): + for config in SDK_CONFIGS: + logger.info("%s:", config["name"]) + releases = get_sdk_releases(config) + if not releases: + logger.info(" no Windows on Snapdragon releases found") + continue + for release in releases: + logger.info(" %s: %s", release["version"], release["name"]) + + +def get_sdk_release(config, version): + version = version or config["default_version"] + version = version.removeprefix("v") + for release in get_sdk_releases(config): + if release["version"] == version: + if not release["sha256"]: + raise RuntimeError(f"{config['name']} {version} does not provide a SHA-256 digest") + return release + raise RuntimeError( + f"No Windows on Snapdragon release for {config['name']} {version}. " + "Run scripts/snapdragon/setup-sdk.py --list-sdk-releases to see available versions." + ) + + +def sha256sum(path): + digest = hashlib.sha256() + with open(path, "rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def download_sdk(release, archive): + while True: + if archive.exists() and sha256sum(archive) == release["sha256"]: + logger.info("Using existing archive %s", archive) + return + + offset = archive.stat().st_size if archive.exists() else 0 + headers = {"User-Agent": "llama.cpp"} + if offset: + headers["Range"] = f"bytes={offset}-" + logger.info("Resuming download of %s at %d MiB", release["name"], offset // (1024 * 1024)) + else: + logger.info("Downloading %s", release["name"]) + + try: + with urlopen(Request(release["url"], headers=headers), timeout=30) as response: + mode = "ab" if offset and response.status == 206 else "wb" + with open(archive, mode) as file: + shutil.copyfileobj(response, file) + except HTTPError as err: + if err.code != 416: + raise RuntimeError(f"Cannot download {release['name']}: {err}") from err + archive.unlink(missing_ok=True) + continue + except (URLError, TimeoutError) as err: + raise RuntimeError(f"Cannot download {release['name']}: {err}") from err + + if sha256sum(archive) == release["sha256"]: + return + raise RuntimeError(f"SHA-256 mismatch for {archive}. Re-run the command to resume the download.") + + +def extract_sdk(config, archive, target_dir): + if not hasattr(tarfile, "data_filter"): + raise RuntimeError("SDK extraction requires Python 3.10.12 or later") + + with tempfile.TemporaryDirectory(prefix=f".{target_dir.name}.tmp-", dir=target_dir.parent) as staging_path: + staging_dir = Path(staging_path) + with tarfile.open(archive, "r:xz") as tar: + tar.extractall(staging_dir, filter=tarfile.data_filter) + + candidates = [staging_dir] + [path for path in staging_dir.iterdir() if path.is_dir()] + extracted_dirs = [path for path in candidates if is_valid_sdk(config, path)] + if len(extracted_dirs) != 1: + raise RuntimeError(f"{config['name']} archive does not contain the expected files") + extracted_dir = extracted_dirs[0] + + backup_dir = None + if target_dir.exists(): + backup_dir = target_dir.parent / f".{target_dir.name}.backup" + if backup_dir.exists(): + raise RuntimeError(f"Cannot replace {target_dir}: backup directory {backup_dir} already exists") + target_dir.replace(backup_dir) + try: + extracted_dir.replace(target_dir) + except Exception: + if backup_dir: + backup_dir.replace(target_dir) + raise + if backup_dir: + shutil.rmtree(backup_dir) + + +def install_sdk(config, version, base_dir, force): + version = (version or config["default_version"]).removeprefix("v") + target_dir = base_dir / config["parent_dir"] / version + if is_valid_sdk(config, target_dir) and not force: + logger.info("Using existing %s at %s", config["name"], target_dir) + return target_dir + + release = get_sdk_release(config, version) + target_dir.parent.mkdir(parents=True, exist_ok=True) + archive = target_dir.parent / release["name"] + download_sdk(release, archive) + logger.info("Extracting %s to %s", config["name"], target_dir) + extract_sdk(config, archive, target_dir) + archive.unlink(missing_ok=True) + return target_dir + + +def set_user_environment(values): + if os.name != "nt": + raise RuntimeError("SDK setup must run on Windows") + + import winreg + + with winreg.CreateKey(winreg.HKEY_CURRENT_USER, "Environment") as key: + for name, value in values.items(): + winreg.SetValueEx(key, name, 0, winreg.REG_SZ, str(value)) + os.environ[name] = str(value) + + import ctypes + + result = ctypes.c_ulong() + ctypes.windll.user32.SendMessageTimeoutW(0xffff, 0x001a, 0, "Environment", 0x0002, 5000, ctypes.byref(result)) + + +def setup_sdks(args): + base_dir = Path(args.sdk_base_dir).expanduser().resolve() + hexagon_config, opencl_config = SDK_CONFIGS + environment = {} + + if args.hexagon is not None: + hexagon_dir = install_sdk(hexagon_config, args.hexagon, base_dir, args.force) + environment["HEXAGON_SDK_ROOT"] = hexagon_dir + environment["HEXAGON_TOOLS_ROOT"] = get_hexagon_tools_dir(hexagon_dir) + if args.opencl is not None: + opencl_dir = install_sdk(opencl_config, args.opencl, base_dir, args.force) + environment["OPENCL_SDK_ROOT"] = opencl_dir + + set_user_environment(environment) + logger.info("SDK environment variables were updated. Start a new terminal before building.") + + +def main(): + logging.basicConfig(level=logging.INFO, format="%(message)s") + parser = argparse.ArgumentParser(description="Install Windows on Snapdragon SDKs for llama.cpp.") + parser.add_argument("--list-sdk-releases", action="store_true", help="List available Windows on Snapdragon SDK releases") + parser.add_argument("--sdk-base-dir", default=DEFAULT_SDK_BASE_DIR, help=r"SDK installation directory (default: C:\Qualcomm)") + parser.add_argument("--hexagon", nargs="?", const=SDK_CONFIGS[0]["default_version"], metavar="VERSION", help="Install the Hexagon SDK, optionally selecting a version") + parser.add_argument("--opencl", nargs="?", const=SDK_CONFIGS[1]["default_version"], metavar="VERSION", help="Install the OpenCL SDK, optionally selecting a version") + parser.add_argument("--force", action="store_true", help="Reinstall selected SDKs even when they already exist") + args = parser.parse_args() + + if args.list_sdk_releases: + if args.sdk_base_dir != DEFAULT_SDK_BASE_DIR or args.hexagon is not None or args.opencl is not None or args.force: + parser.error("Installation options cannot be combined with --list-sdk-releases") + list_sdk_releases() + return + if args.hexagon is None and args.opencl is None: + parser.error("Select at least one SDK with --hexagon or --opencl") + if os.name != "nt": + parser.error("SDK setup must run on Windows") + setup_sdks(args) + + +if __name__ == "__main__": + try: + main() + except KeyboardInterrupt: + logger.info("\nInterrupted by user.") + sys.exit(130) + except RuntimeError as err: + logger.error("Error: %s", err) + sys.exit(1) diff --git a/scripts/snapdragon/windows/run-bench.ps1 b/scripts/snapdragon/windows/run-bench.ps1 deleted file mode 100644 index 6eb656e66d3..00000000000 --- a/scripts/snapdragon/windows/run-bench.ps1 +++ /dev/null @@ -1,48 +0,0 @@ - -#!/usr/bin/env pwsh - -# Basedir on device -$basedir=".\pkg-snapdragon" - -$cli_opts=$args - -$model="Llama-3.2-3B-Instruct-Q4_0.gguf" -if ($null -ne $env:M) { - $model=$env:M -} - -$device="HTP0" -if ($null -ne $env:D) { - $device=$env:D -} - -if ($null -ne $env:V) { - $env:GGML_HEXAGON_VERBOSE=$env:V -} - -if ($null -ne $env:PROF) { - $env:GGML_HEXAGON_PROFILE=$env:PROF -} - -if ($null -ne $env:OPSTAGE) { - $env:GGML_HEXAGON_OPSTAGE=$env:OPSTAGE -} - -if ($null -ne $env:NHVX) { - $env:GGML_HEXAGON_NHVX=$env:NHVX -} - -if ($null -ne $env:NDEV) { - $env:GGML_HEXAGON_NDEV=$env:NDEV -} - -if ($null -ne $env:HB) { - $env:GGML_HEXAGON_HOSTBUF=$env:HB -} - -$env:ADSP_LIBRARY_PATH="$basedir\lib" - -& "$basedir\bin\llama-bench.exe" ` - --load-mode none -m $basedir\..\..\gguf\$model ` - --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 ` - --ubatch-size 1024 -ngl 99 --device $device $cli_opts diff --git a/scripts/snapdragon/windows/run-cli.ps1 b/scripts/snapdragon/windows/run-cli.ps1 deleted file mode 100644 index 5da8bff33e3..00000000000 --- a/scripts/snapdragon/windows/run-cli.ps1 +++ /dev/null @@ -1,53 +0,0 @@ - -#!/usr/bin/env pwsh - -# Basedir on device -$basedir=".\pkg-snapdragon" - -$cli_opts=$args - -$model="Llama-3.2-3B-Instruct-Q4_0.gguf" -if ($null -ne $env:M) { - $model=$env:M -} - -$device="HTP0" -if ($null -ne $env:D) { - $device=$env:D -} - -if ($null -ne $env:V) { - $env:GGML_HEXAGON_VERBOSE=$env:V -} - -if ($null -ne $env:SCHED) { - $env:GGML_SCHED_DEBUG=$env:SCHED; $cli_opts="$cli_opts -v" -} - -if ($null -ne $env:PROF) { - $env:GGML_HEXAGON_PROFILE=$env:PROF -} - -if ($null -ne $env:OPSTAGE) { - $env:GGML_HEXAGON_OPSTAGE=$env:OPSTAGE -} - -if ($null -ne $env:NHVX) { - $env:GGML_HEXAGON_NHVX=$env:NHVX -} - -if ($null -ne $env:NDEV) { - $env:GGML_HEXAGON_NDEV=$env:NDEV -} - -if ($null -ne $env:HB) { - $env:GGML_HEXAGON_HOSTBUF=$env:HB -} - -$env:ADSP_LIBRARY_PATH="$basedir\lib" - -& "$basedir\bin\llama-cli.exe" ` - --load-mode none -m $basedir\..\..\gguf\$model ` - --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 ` - --ctx-size 8192 --ubatch-size 1024 -fa on ` - -ngl 99 --device $device $cli_opts diff --git a/scripts/snapdragon/windows/run-completion.ps1 b/scripts/snapdragon/windows/run-completion.ps1 deleted file mode 100644 index 08ef139b7e2..00000000000 --- a/scripts/snapdragon/windows/run-completion.ps1 +++ /dev/null @@ -1,53 +0,0 @@ - -#!/usr/bin/env pwsh - -# Basedir on device -$basedir=".\pkg-snapdragon" - -$cli_opts=$args - -$model="Llama-3.2-3B-Instruct-Q4_0.gguf" -if ($null -ne $env:M) { - $model=$env:M -} - -$device="HTP0" -if ($null -ne $env:D) { - $device=$env:D -} - -if ($null -ne $env:V) { - $env:GGML_HEXAGON_VERBOSE=$env:V -} - -if ($null -ne $env:SCHED) { - $env:GGML_SCHED_DEBUG=$env:SCHED; $cli_opts="$cli_opts -v" -} - -if ($null -ne $env:PROF) { - $env:GGML_HEXAGON_PROFILE=$env:PROF -} - -if ($null -ne $env:OPSTAGE) { - $env:GGML_HEXAGON_OPSTAGE=$env:OPSTAGE -} - -if ($null -ne $env:NHVX) { - $env:GGML_HEXAGON_NHVX=$env:NHVX -} - -if ($null -ne $env:NDEV) { - $env:GGML_HEXAGON_NDEV=$env:NDEV -} - -if ($null -ne $env:HB) { - $env:GGML_HEXAGON_HOSTBUF=$env:HB -} - -$env:ADSP_LIBRARY_PATH="$basedir\lib" - -& "$basedir\bin\llama-completion.exe" ` - --load-mode none -m $basedir\..\..\gguf\$model ` - --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 ` - --ctx-size 8192 --ubatch-size 1024 -fa on ` - -ngl 99 -no-cnv --device $device $cli_opts diff --git a/scripts/snapdragon/windows/run-mtmd.ps1 b/scripts/snapdragon/windows/run-mtmd.ps1 deleted file mode 100644 index 6e270ec90b5..00000000000 --- a/scripts/snapdragon/windows/run-mtmd.ps1 +++ /dev/null @@ -1,68 +0,0 @@ -#!/usr/bin/env pwsh - -# Basedir on device -$basedir=".\pkg-snapdragon" - -$cli_opts=$args - -$model="gemma-3-4b-it-Q4_0.gguf" -if ($null -ne $env:M) { - $model=$env:M -} - -$mmproj="mmproj-F16.gguf" -if ($null -ne $env:MMPROJ) { - $mmproj=$env:MMPROJ -} - -$image="" -if ($null -ne $env:IMG) { - $image=$env:IMG -} - -$device="HTP0" -if ($null -ne $env:D) { - $device=$env:D -} - -if ($null -ne $env:V) { - $env:GGML_HEXAGON_VERBOSE=$env:V -} - -if ($null -ne $env:SCHED) { - $env:GGML_SCHED_DEBUG=$env:SCHED; $cli_opts="$cli_opts -v" -} - -if ($null -ne $env:PROF) { - $env:GGML_HEXAGON_PROFILE=$env:PROF -} - -if ($null -ne $env:OPSTAGE) { - $env:GGML_HEXAGON_OPSTAGE=$env:OPSTAGE -} - -if ($null -ne $env:NHVX) { - $env:GGML_HEXAGON_NHVX=$env:NHVX -} - -if ($null -ne $env:NDEV) { - $env:GGML_HEXAGON_NDEV=$env:NDEV -} - -if ($null -ne $env:HB) { - $env:GGML_HEXAGON_HOSTBUF=$env:HB -} - -if ($null -ne $env:MTMD_DEVICE) { - $env:MTMD_BACKEND_DEVICE=$env:MTMD_DEVICE -} - -$env:ADSP_LIBRARY_PATH="$basedir\lib" - -& "$basedir\bin\llama-mtmd-cli.exe" ` - --load-mode none -m $basedir\..\..\gguf\$model ` - --mmproj $basedir\..\..\gguf\$mmproj ` - --image $basedir\..\..\gguf\$image ` - --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 ` - --ctx-size 8192 --ubatch-size 1024 -fa on ` - -ngl 99 --device $device -v $cli_opts diff --git a/scripts/snapdragon/windows/run-tool.ps1 b/scripts/snapdragon/windows/run-tool.ps1 deleted file mode 100644 index 15c880f2dbd..00000000000 --- a/scripts/snapdragon/windows/run-tool.ps1 +++ /dev/null @@ -1,56 +0,0 @@ - -#!/usr/bin/env pwsh - -# Basedir on device -$basedir=".\pkg-snapdragon" - -if ($args.Count -eq 0) { - Write-Host "No arguments provided.Expected the tool and argument to run." - exit -1 -} - -$tool=$args[0] -$cli_opts=@() - -if ($args.Count -gt 1) { - $cli_opts=$args[1..($args.Count - 1)] - $remainingArgs = $args[1..($args.Count - 1)] -} - -$device="HTP0" -if ($null -ne $env:D) { - $device=$env:D -} - -if ($null -ne $env:V) { - $env:GGML_HEXAGON_VERBOSE=$env:V -} - -if ($null -ne $env:SCHED) { - $env:GGML_SCHED_DEBUG=$env:SCHED; $cli_opts="$cli_opts -v" -} - -if ($null -ne $env:PROF) { - $env:GGML_HEXAGON_PROFILE=$env:PROF -} - -if ($null -ne $env:OPSTAGE) { - $env:GGML_HEXAGON_OPSTAGE=$env:OPSTAGE -} - -if ($null -ne $env:NHVX) { - $env:GGML_HEXAGON_NHVX=$env:NHVX -} - -if ($null -ne $env:NDEV) { - $env:GGML_HEXAGON_NDEV=$env:NDEV -} - -if ($null -ne $env:HB) { - $env:GGML_HEXAGON_HOSTBUF=$env:HB -} - -$env:ADSP_LIBRARY_PATH="$basedir\lib" - -& "$basedir\bin\$tool" ` - $cli_opts diff --git a/scripts/snapdragon/windows/setup-build.ps1 b/scripts/snapdragon/windows/setup-build.ps1 deleted file mode 100644 index d8ef24d4413..00000000000 --- a/scripts/snapdragon/windows/setup-build.ps1 +++ /dev/null @@ -1,105 +0,0 @@ -# Requires Run as Administrator is NOT strictly necessary for User-scope env vars, -# but recommended for creating directories in C:\ root if permissions are restricted. - -$ErrorActionPreference = "Stop" - -# --- Configuration --- -$BaseDir = "C:\Qualcomm" - -# SDK 1: Hexagon -$HexagonUrl = "https://github.com/snapdragon-toolchain/hexagon-sdk/releases/download/v6.6.0.0/hexagon-sdk-v6.6.0.0-arm64-wos.tar.xz" -$HexagonParent = Join-Path $BaseDir "Hexagon_SDK" -$HexagonSdkVersion = "6.6.0.0" -$HexagonToolsVersion = "19.0.07" -$HexagonSdkTarget = Join-Path $HexagonParent $HexagonSdkVersion -$HexagonToolsTarget = Join-Path $HexagonSdkTarget "\tools\HEXAGON_Tools\$HexagonToolsVersion" - -# SDK 2: OpenCL -$OpenCLUrl = "https://github.com/snapdragon-toolchain/opencl-sdk/releases/download/v2.3.2/adreno-opencl-sdk-v2.3.2-arm64-wos.tar.xz" -$OpenCLParent = Join-Path $BaseDir "OpenCL_SDK" -$OpenCLVersion = "2.3.2" -$OpenCLTarget = Join-Path $OpenCLParent $OpenCLVersion - -# --- Helper Function --- -function Install-QualcommSDK { - param ( - [string]$Url, - [string]$ParentDir, - [string]$TargetDir, - [string]$Name - ) - - # 1. Create Parent Directory - if (-not (Test-Path -Path $ParentDir)) { - Write-Host "Creating directory: $ParentDir" -ForegroundColor Cyan - New-Item -Path $ParentDir -ItemType Directory -Force | Out-Null - } - - # 2. Check for Specific Version Directory - if (Test-Path -Path $TargetDir) { - Write-Host "$Name ($TargetDir) already exists. Skipping download." -ForegroundColor Green - } - else { - Write-Host "$Name not found. preparing to download..." -ForegroundColor Yellow - - # Create the target directory to extract into - New-Item -Path $TargetDir -ItemType Directory -Force | Out-Null - - # Define temporary archive path - $TempFile = Join-Path $ParentDir "temp_sdk.tar.xz" - - try { - # Download - Write-Host "Downloading from: $Url" - Invoke-WebRequest -Uri $Url -OutFile $TempFile - - # Untar - # Note: We assume Windows includes tar.exe (Win 10 build 17063+) - Write-Host "Extracting archive to $TargetDir..." - - # We use -C to extract contents INTO the target directory created above - tar -xJvf $TempFile -C $TargetDir\.. - - Write-Host "Extraction complete." -ForegroundColor Green - } - catch { - Write-Error "Failed to download or extract $Name. Error: $_" - # Cleanup target dir if failed so script tries again next time - Remove-Item -Path $TargetDir -Recurse -Force -ErrorAction SilentlyContinue - } - finally { - # Cleanup Archive - if (Test-Path $TempFile) { Remove-Item $TempFile -Force } - } - } -} - -# --- Execution --- - -# 1. Ensure Base C:\Qualcomm exists -if (-not (Test-Path $BaseDir)) { - New-Item -Path $BaseDir -ItemType Directory -Force | Out-Null -} - -# 2. Run Install Logic -Install-QualcommSDK -Url $HexagonUrl -ParentDir $HexagonParent -TargetDir $HexagonSdkTarget -Name "Hexagon SDK" -Install-QualcommSDK -Url $OpenCLUrl -ParentDir $OpenCLParent -TargetDir $OpenCLTarget -Name "OpenCL SDK" - -# --- Environment Variables --- - -Write-Host "`nSetting Environment Variables..." -ForegroundColor Cyan - -# Set OPENCL_SDK_ROOT -[System.Environment]::SetEnvironmentVariable('OPENCL_SDK_ROOT', $OpenCLTarget, [System.EnvironmentVariableTarget]::User) -$env:OPENCL_SDK_ROOT = $OpenCLTarget # Set for current session as well -Write-Host "OPENCL_SDK_ROOT set to: $OpenCLTarget" - -# Set HEXAGON_SDK_ROOT -[System.Environment]::SetEnvironmentVariable('HEXAGON_SDK_ROOT', $HexagonSdkTarget, [System.EnvironmentVariableTarget]::User) -$env:HEXAGON_SDK_ROOT = $HexagonSdkTarget # Set for current session as well -Write-Host "HEXAGON_SDK_ROOT set to: $HexagonSdkTarget" - -# Set HEXAGON_SDK_ROOT -[System.Environment]::SetEnvironmentVariable('HEXAGON_TOOLS_ROOT', $HexagonToolsTarget, [System.EnvironmentVariableTarget]::User) -$env:HEXAGON_TOOLS_ROOT = $HexagonToolsTarget # Set for current session as well -Write-Host "HEXAGON_TOOLS_ROOT set to: $HexagonToolsTarget" diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index c6df19f2ecf..8922dc12adc 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -31,6 +31,7 @@ add_library(llama llama-memory.cpp llama-memory-hybrid.cpp llama-memory-hybrid-iswa.cpp + llama-memory-hybrid-idx.cpp llama-memory-recurrent.cpp llama-mmap.cpp llama-model-loader.cpp diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index eecf444fcf3..5e61f61f7f0 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -40,6 +40,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_QWEN3VLMOE, "qwen3vlmoe" }, { LLM_ARCH_QWEN35, "qwen35" }, { LLM_ARCH_QWEN35MOE, "qwen35moe" }, + { LLM_ARCH_QWEN4EXP, "qwen4exp" }, { LLM_ARCH_PHI2, "phi2" }, { LLM_ARCH_PHI3, "phi3" }, { LLM_ARCH_PHIMOE, "phimoe" }, @@ -293,6 +294,17 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_HYPER_CONNECTION_COUNT, "%s.hyper_connection.count" }, { LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, "%s.hyper_connection.sinkhorn_iterations" }, { LLM_KV_HYPER_CONNECTION_EPSILON, "%s.hyper_connection.epsilon" }, + { LLM_KV_HYPER_CONNECTION_LOW_RANK, "%s.hyper_connection.low_rank" }, + + { LLM_KV_PLE_LAYERS, "%s.ple.layers" }, + { LLM_KV_PLE_NGRAM_SIZE, "%s.ple.ngram_size" }, + { LLM_KV_PLE_HEADS_PER_NGRAM, "%s.ple.heads_per_ngram" }, + { LLM_KV_PLE_CONV_KERNEL, "%s.ple.conv_kernel" }, + { LLM_KV_PLE_LAYER_MULTIPLIERS, "%s.ple.layer_multipliers" }, + { LLM_KV_PLE_HEAD_OFFSETS, "%s.ple.head_offsets" }, + { LLM_KV_PLE_HEAD_VOCAB_SIZES, "%s.ple.head_vocab_sizes" }, + { LLM_KV_PLE_EOS_TOKEN_ID, "%s.ple.eos_token_id" }, + { LLM_KV_PLE_IMAGE_TOKEN_ID, "%s.ple.image_token_id" }, { LLM_KV_HASH_LAYER_COUNT, "%s.hash_layer_count" }, @@ -344,6 +356,12 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_NORM_BEFORE_RESIDUAL, "%s.norm_before_residual" }, { LLM_KV_NORM_BEFORE_FC, "%s.norm_before_fc" }, + { LLM_KV_DFLASH_BLOCK_SIZE, "%s.block_size" }, + { LLM_KV_DFLASH_CONV_KERNEL_SIZE, "%s.conv_kernel_size" }, + { LLM_KV_DFLASH_CONV_GROUP_SIZE, "%s.conv_group_size" }, + { LLM_KV_DFLASH_SELECTOR_RANK, "%s.selector_rank" }, + { LLM_KV_DFLASH_SELECTOR_TOP_K, "%s.selector_top_k" }, + { LLM_KV_SHORTCONV_L_CACHE, "%s.shortconv.l_cache" }, // sentence-transformers dense modules feature dims { LLM_KV_DENSE_2_FEAT_IN, "%s.dense_2_feat_in" }, @@ -500,12 +518,29 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_HC_HEAD_FN, "output_hc_fn" }, { LLM_TENSOR_HC_HEAD_BASE, "output_hc_base" }, { LLM_TENSOR_HC_HEAD_SCALE, "output_hc_scale" }, + { LLM_TENSOR_HC_HEAD_NORM, "output_hc_norm" }, + { LLM_TENSOR_HC_HEAD_DOWN, "output_hc_down" }, + { LLM_TENSOR_HC_HEAD_UP, "output_hc_up" }, { LLM_TENSOR_HC_ATTN_FN, "blk.%d.hc_attn_fn" }, { LLM_TENSOR_HC_ATTN_BASE, "blk.%d.hc_attn_base" }, { LLM_TENSOR_HC_ATTN_SCALE, "blk.%d.hc_attn_scale" }, { LLM_TENSOR_HC_FFN_FN, "blk.%d.hc_ffn_fn" }, { LLM_TENSOR_HC_FFN_BASE, "blk.%d.hc_ffn_base" }, { LLM_TENSOR_HC_FFN_SCALE, "blk.%d.hc_ffn_scale" }, + { LLM_TENSOR_HC_ATTN_NORM, "blk.%d.hc_attn_norm" }, + { LLM_TENSOR_HC_ATTN_DOWN, "blk.%d.hc_attn_down" }, + { LLM_TENSOR_HC_ATTN_UP, "blk.%d.hc_attn_up" }, + { LLM_TENSOR_HC_ATTN_INJECT, "blk.%d.hc_attn_inject" }, + { LLM_TENSOR_HC_FFN_NORM, "blk.%d.hc_ffn_norm" }, + { LLM_TENSOR_HC_FFN_DOWN, "blk.%d.hc_ffn_down" }, + { LLM_TENSOR_HC_FFN_UP, "blk.%d.hc_ffn_up" }, + { LLM_TENSOR_HC_FFN_INJECT, "blk.%d.hc_ffn_inject" }, + { LLM_TENSOR_PLE_KEY, "blk.%d.ple_key" }, + { LLM_TENSOR_PLE_VALUE, "blk.%d.ple_value" }, + { LLM_TENSOR_PLE_NORM_KEY, "blk.%d.ple_norm_key" }, + { LLM_TENSOR_PLE_NORM_QUERY, "blk.%d.ple_norm_query" }, + { LLM_TENSOR_PLE_NORM_CONV, "blk.%d.ple_norm_conv" }, + { LLM_TENSOR_PLE_CONV1D, "blk.%d.ple_conv1d" }, { LLM_TENSOR_ATTN_COMPRESSOR_WKV, "blk.%d.attn_compressor_kv" }, { LLM_TENSOR_ATTN_COMPRESSOR_WGATE, "blk.%d.attn_compressor_gate" }, { LLM_TENSOR_ATTN_COMPRESSOR_APE, "blk.%d.attn_compressor_ape" }, @@ -651,6 +686,13 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_DSPARK_MARKOV_W1, "markov_w1" }, { LLM_TENSOR_DSPARK_MARKOV_W2, "markov_w2" }, { LLM_TENSOR_DSPARK_CONF_PROJ, "conf_proj" }, + { LLM_TENSOR_DFLASH_ATTN_CONV_BASE, "blk.%d.attn_conv_base" }, + { LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, "blk.%d.attn_conv_proj" }, + { LLM_TENSOR_DFLASH_FFN_CONV_BASE, "blk.%d.ffn_conv_base" }, + { LLM_TENSOR_DFLASH_FFN_CONV_PROJ, "blk.%d.ffn_conv_proj" }, + { LLM_TENSOR_DFLASH_SELECTOR_PREV, "selector_predecessor" }, + { LLM_TENSOR_DFLASH_SELECTOR_NEXT, "selector_successor" }, + { LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, "selector_hidden" }, }; // declare information about the model weight tensors: @@ -704,12 +746,29 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_HC_HEAD_FN, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, {LLM_TENSOR_HC_HEAD_BASE, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_ADD}}, {LLM_TENSOR_HC_HEAD_SCALE, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, + {LLM_TENSOR_HC_HEAD_NORM, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, + {LLM_TENSOR_HC_HEAD_DOWN, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_HEAD_UP, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, {LLM_TENSOR_HC_ATTN_FN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_HC_ATTN_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}}, {LLM_TENSOR_HC_ATTN_SCALE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_HC_FFN_FN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_HC_FFN_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}}, {LLM_TENSOR_HC_FFN_SCALE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_HC_ATTN_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_HC_ATTN_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_ATTN_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_ATTN_INJECT, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_FFN_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_HC_FFN_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_FFN_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_FFN_INJECT, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_PLE_KEY, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_PLE_VALUE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_PLE_NORM_KEY, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_PLE_NORM_QUERY, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_PLE_NORM_CONV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_PLE_CONV1D, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SSM_CONV}}, {LLM_TENSOR_ATTN_COMPRESSOR_WKV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_COMPRESSOR_WGATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_COMPRESSOR_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}}, @@ -916,6 +975,13 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_DSPARK_MARKOV_W1, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, {LLM_TENSOR_DSPARK_MARKOV_W2, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, {LLM_TENSOR_DSPARK_CONF_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_DFLASH_ATTN_CONV_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_DFLASH_FFN_CONV_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_DFLASH_FFN_CONV_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_DFLASH_SELECTOR_PREV, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, + {LLM_TENSOR_DFLASH_SELECTOR_NEXT, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, + {LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, }; LLM_KV::LLM_KV(llm_arch arch, const char * suffix) : arch(arch), suffix(suffix) {} @@ -1009,6 +1075,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_KIMI_K3: case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: + case LLM_ARCH_QWEN4EXP: case LLM_ARCH_DEEPSEEK4: case LLM_ARCH_MINIMAX_01: return true; diff --git a/src/llama-arch.h b/src/llama-arch.h index 7159e23bf7a..ca7d55a5fd7 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -45,6 +45,7 @@ enum llm_arch { LLM_ARCH_QWEN3VLMOE, LLM_ARCH_QWEN35, LLM_ARCH_QWEN35MOE, + LLM_ARCH_QWEN4EXP, LLM_ARCH_PHI2, LLM_ARCH_PHI3, LLM_ARCH_PHIMOE, @@ -298,6 +299,17 @@ enum llm_kv { LLM_KV_HYPER_CONNECTION_COUNT, LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, LLM_KV_HYPER_CONNECTION_EPSILON, + LLM_KV_HYPER_CONNECTION_LOW_RANK, + + LLM_KV_PLE_LAYERS, + LLM_KV_PLE_NGRAM_SIZE, + LLM_KV_PLE_HEADS_PER_NGRAM, + LLM_KV_PLE_CONV_KERNEL, + LLM_KV_PLE_LAYER_MULTIPLIERS, + LLM_KV_PLE_HEAD_OFFSETS, + LLM_KV_PLE_HEAD_VOCAB_SIZES, + LLM_KV_PLE_EOS_TOKEN_ID, + LLM_KV_PLE_IMAGE_TOKEN_ID, LLM_KV_HASH_LAYER_COUNT, @@ -387,6 +399,11 @@ enum llm_kv { LLM_KV_TARGET_LAYERS, LLM_KV_TARGET_HIDDEN_SIZE, + LLM_KV_DFLASH_BLOCK_SIZE, + LLM_KV_DFLASH_CONV_KERNEL_SIZE, + LLM_KV_DFLASH_CONV_GROUP_SIZE, + LLM_KV_DFLASH_SELECTOR_RANK, + LLM_KV_DFLASH_SELECTOR_TOP_K, LLM_KV_NORM_BEFORE_RESIDUAL, LLM_KV_NORM_BEFORE_FC, @@ -565,12 +582,29 @@ enum llm_tensor { LLM_TENSOR_HC_HEAD_FN, LLM_TENSOR_HC_HEAD_BASE, LLM_TENSOR_HC_HEAD_SCALE, + LLM_TENSOR_HC_HEAD_NORM, // qwen4exp + LLM_TENSOR_HC_HEAD_DOWN, // qwen4exp + LLM_TENSOR_HC_HEAD_UP, // qwen4exp LLM_TENSOR_HC_ATTN_FN, LLM_TENSOR_HC_ATTN_BASE, LLM_TENSOR_HC_ATTN_SCALE, LLM_TENSOR_HC_FFN_FN, LLM_TENSOR_HC_FFN_BASE, LLM_TENSOR_HC_FFN_SCALE, + LLM_TENSOR_HC_ATTN_NORM, // qwen4exp + LLM_TENSOR_HC_ATTN_DOWN, // qwen4exp + LLM_TENSOR_HC_ATTN_UP, // qwen4exp + LLM_TENSOR_HC_ATTN_INJECT, // qwen4exp + LLM_TENSOR_HC_FFN_NORM, // qwen4exp + LLM_TENSOR_HC_FFN_DOWN, // qwen4exp + LLM_TENSOR_HC_FFN_UP, // qwen4exp + LLM_TENSOR_HC_FFN_INJECT, // qwen4exp + LLM_TENSOR_PLE_KEY, // qwen4exp + LLM_TENSOR_PLE_VALUE, // qwen4exp + LLM_TENSOR_PLE_NORM_KEY, // qwen4exp + LLM_TENSOR_PLE_NORM_QUERY, // qwen4exp + LLM_TENSOR_PLE_NORM_CONV, // qwen4exp + LLM_TENSOR_PLE_CONV1D, // qwen4exp LLM_TENSOR_ATTN_COMPRESSOR_WKV, LLM_TENSOR_ATTN_COMPRESSOR_WGATE, LLM_TENSOR_ATTN_COMPRESSOR_APE, @@ -659,6 +693,13 @@ enum llm_tensor { LLM_TENSOR_DSPARK_MARKOV_W1, LLM_TENSOR_DSPARK_MARKOV_W2, LLM_TENSOR_DSPARK_CONF_PROJ, + LLM_TENSOR_DFLASH_ATTN_CONV_BASE, + LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, + LLM_TENSOR_DFLASH_FFN_CONV_BASE, + LLM_TENSOR_DFLASH_FFN_CONV_PROJ, + LLM_TENSOR_DFLASH_SELECTOR_PREV, + LLM_TENSOR_DFLASH_SELECTOR_NEXT, + LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, }; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0402044da6b..9aed8013327 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -231,10 +231,10 @@ llama_context::llama_context( cparams.fused_gdn_ar = true; cparams.fused_gdn_ch = true; - cparams.auto_fgdn = true; + cparams.auto_fgdn = false; - cparams.fused_lid = true; - cparams.auto_flid = true; + cparams.fused_lid = true; + cparams.auto_flid = false; cparams.fused_dsv4_hc_pre = true; cparams.fused_dsv4_hc_comb = true; @@ -661,11 +661,19 @@ void llama_context::sched_reserve() { // reserve again with pp graph to avoid ggml-alloc reallocations during inference { - // TODO: not sure if the following graph would be worst case for multi-stream KV caches: - // - // auto * gf = graph_reserve(n_tokens, 1, n_tokens, mctx.get()); - // - auto * gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(), model.hparams.no_alloc); + // TODO: the worst case graph is not always reached for `n_seqs > 1` + // need to implement a more robust mechanism that tries a few different inputs and analyzes the results + ggml_cgraph * gf = nullptr; + switch (model.arch) { + case LLM_ARCH_MINIMAX_01: + // the `inp_diag_decay` tensor size scales with `n_seq_tokens^2` which + // makes `n_seqs == 1` use more memory for the compute graph compared to `n_seqs > 1` + gf = graph_reserve(n_tokens, 1, n_outputs_pp, mctx.get(), model.hparams.no_alloc); + break; + default: + gf = graph_reserve(n_tokens, n_seqs, n_outputs_pp, mctx.get(), model.hparams.no_alloc); + }; + if (!gf) { throw std::runtime_error("failed to allocate compute pp buffers"); } @@ -2301,12 +2309,17 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { model.arch == LLM_ARCH_BAILINGMOE3 || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || + model.arch == LLM_ARCH_QWEN4EXP || model.arch == LLM_ARCH_DEEPSEEK4 || (model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) || model.arch == LLM_ARCH_NANBEIGE || model.arch == LLM_ARCH_MINIMAX_01 || model.arch == LLM_ARCH_MINIMAX_M3) { res = std::max(n_tokens * 40, 32u * model.n_tensors()); + } else if (model.arch == LLM_ARCH_DFLASH && model.hparams.dflash_selector_rank > 0) { + // DFlash2's convolutions and selector are shape work rather than matmuls, + // so they cost ~8.6 nodes per tensor against ~5.9 for a plain DFlash draft + res = std::max(1024u, 12u*model.n_tensors()); } else { res = std::max(1024u, 8u*model.n_tensors()); for (const auto & lora : model.loras) { @@ -2887,13 +2900,83 @@ class llama_io_read_device : public llama_io_read_i { for (auto & [buft, mbuf] : mbufs_new) { const auto & mbuf_cur = mbufs.at(buft); - if (!mbuf_cur.buf || mbuf_cur.n_tensors != mbuf.n_tensors || mbuf_cur.total_size != mbuf.total_size) { + if (!mbuf_cur.buf || mbuf_cur.total_size != mbuf.total_size) { GGML_ABORT("%s: memory buffer mismatch\n", __func__); } - for (size_t i = 0; i < mbuf_cur.org.size(); ++i) { - ggml_backend_tensor_copy(mbuf_cur.cpy[i], mbuf.org[i]); + if (mbuf_cur.n_tensors == mbuf.n_tensors) { + // same chunking: copy 1:1 by index + for (size_t i = 0; i < mbuf_cur.org.size(); ++i) { + GGML_ASSERT(ggml_nbytes(mbuf_cur.cpy[i]) == ggml_nbytes(mbuf.org[i])); + ggml_backend_tensor_copy(mbuf_cur.cpy[i], mbuf.org[i]); + } + continue; } + + // different chunking: copy the write-side data (mbuf_cur.cpy) into the read-side targets (mbuf.org) + // with a byte cursor. Write and read enumerate the same logical data in the same order but may chunk + // it differently, so copy across tensor boundaries rather than 1:1 by index. + const size_t total = mbuf_cur.total_size; + + ggml_init_params params_scratch = { + /*.mem_size =*/ 2*(mbuf_cur.cpy.size() + mbuf.org.size())*ggml_tensor_overhead(), + /*.mem_buffer =*/ NULL, + /*.no_alloc =*/ true, + }; + ggml_context * ctx_scratch = ggml_init(params_scratch); + + size_t src_pos = 0; + size_t dst_pos = 0; + size_t src_j = 0; + size_t dst_i = 0; + size_t src_base = 0; + size_t dst_base = 0; + + while (src_pos < total) { + const auto & src_t = mbuf_cur.cpy[src_j]; + const auto & dst_t = mbuf.org[dst_i]; + + const size_t src_size = ggml_nbytes(src_t); + const size_t dst_size = ggml_nbytes(dst_t); + + const size_t src_off = src_pos - src_base; + const size_t dst_off = dst_pos - dst_base; + + const size_t n_copy = std::min(src_size - src_off, dst_size - dst_off); + + const size_t el = ggml_element_size(src_t); + const int64_t n_el = (int64_t) (n_copy / el); + + auto * src_v = ggml_view_1d(ctx_scratch, src_t, n_el, src_off); + ggml_backend_view_init(src_v); + auto * dst_v = ggml_view_1d(ctx_scratch, dst_t, n_el, dst_off); + ggml_backend_view_init(dst_v); + + ggml_backend_tensor_copy(src_v, dst_v); + + src_pos += n_copy; + dst_pos += n_copy; + + if (src_pos - src_base == src_size) { + src_base = src_pos; + ++src_j; + } + if (dst_pos - dst_base == dst_size) { + dst_base = dst_pos; + ++dst_i; + } + } + + GGML_ASSERT(src_pos == total && dst_pos == total); + // any tensors left unvisited hold no data + for (size_t i = src_j; i < mbuf_cur.cpy.size(); ++i) { + GGML_ASSERT(ggml_nbytes(mbuf_cur.cpy[i]) == 0); + } + for (size_t i = dst_i; i < mbuf.org.size(); ++i) { + GGML_ASSERT(ggml_nbytes(mbuf.org[i]) == 0); + } + + ggml_free(ctx_scratch); } GGML_ASSERT(buf_size == 0); diff --git a/src/llama-ext.h b/src/llama-ext.h index 35d6e58adfa..92a759b7a0a 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -120,6 +120,8 @@ LLAMA_API llama_context * llama_get_ctx_other(struct llama_context * ctx); // model/context data extraction // +LLAMA_API int32_t llama_model_dflash_selector_top_k(const struct llama_model * model); + // returns pointer to the target-model layer indices LLAMA_API const int32_t * llama_model_target_layer_ids (const struct llama_model * model); // returns the number of extracted layers from target model diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index cbe31134ff4..6a820c61c99 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -201,7 +201,11 @@ uint32_t llama_hparams::n_embd_r() const { // TODO: maybe support other convolution strides than 1 // NOTE: since the first column of the conv_state is shifted out each time, it's not actually needed // Corresponds to Mamba's conv_states size - return (ssm_d_conv > 0 ? ssm_d_conv - 1 : 0) * (ssm_d_inner + 2*ssm_n_group*ssm_d_state); + const uint32_t n_conv = (ssm_d_conv > 0 ? ssm_d_conv - 1 : 0) * (ssm_d_inner + 2*ssm_n_group*ssm_d_state); + + // PLE conv history needs its own row: Meta splits cache_r_l by head, so a history packed behind the first is unaddressable + // it lives in cache_ple_r_l instead, mirrored like the rest of the PLE module + return n_conv; } uint32_t llama_hparams::n_embd_s() const { @@ -236,6 +240,23 @@ bool llama_hparams::is_recr(uint32_t il) const { GGML_ABORT("%s: il (%u) out of bounds (n_layer_all: %u)\n", __func__, il, n_layer_all); } +uint32_t llama_hparams::ple_conv_state() const { + if (ple_n_heads == 0 || ple_conv_kernel == 0) { + return 0; + } + + // dilation equals the n-gram size, matching the reference module + return (ple_conv_kernel - 1) * ple_ngram_size * dsv4_hc_mult * n_embd; +} + +bool llama_hparams::is_ple(uint32_t il) const { + if (il < n_layer_all) { + return is_ple_impl[il]; + } + + GGML_ABORT("%s: il (%u) out of bounds (n_layer_all: %u)\n", __func__, il, n_layer_all); +} + uint32_t llama_hparams::n_pos_per_embd() const { return rope_type == LLAMA_ROPE_TYPE_MROPE || rope_type == LLAMA_ROPE_TYPE_IMROPE ? 4 : 1; } diff --git a/src/llama-hparams.h b/src/llama-hparams.h index c3c14292c32..1411692a890 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -3,12 +3,15 @@ #include "llama.h" #include +#include #include #include // bump if necessary #define LLAMA_MAX_LAYERS 512 #define LLAMA_MAX_EXPERTS 1024 // Kimi K3 +#define LLAMA_MAX_PLE_NGRAM 8 // qwen4exp +#define LLAMA_MAX_PLE_HEADS 64 // qwen4exp enum llama_expert_gating_func_type { LLAMA_EXPERT_GATING_FUNC_TYPE_NONE = 0, @@ -223,6 +226,12 @@ struct llama_hparams { // output embedding dimension (0 = use n_embd) uint32_t n_embd_out_impl = 0; + uint32_t dflash_block_size = 0; + uint32_t dflash_conv_kernel_size = 0; + uint32_t dflash_conv_group_size = 0; + uint32_t dflash_selector_rank = 0; + uint32_t dflash_selector_top_k = 0; + // llama4 smallthinker uint32_t n_moe_layer_step = 0; uint32_t n_no_rope_layer_step = 4; @@ -270,6 +279,30 @@ struct llama_hparams { float dsv4_hc_eps = 0.0f; std::array dsv4_compress_ratios; + // 0 = full rank (DeepSeek-V4) + uint32_t hc_low_rank = 0; + + uint32_t ple_ngram_size = 0; + uint32_t ple_heads_per_ngram = 0; + uint32_t ple_conv_kernel = 0; + uint32_t ple_n_heads = 0; // (ngram_size - 1) * heads_per_ngram + uint32_t ple_head_dim = 0; + uint32_t ple_eos_token_id = 0; + // the id the PLE hash stands in at image positions; 0 makes the loader fall back to EOS + uint32_t ple_image_token_id = 0; + // the file lists PLE layer indices, so this is never a per-layer gguf array and can hold one bit per layer + std::bitset is_ple_impl; + // the hash multipliers reach ~2e13 and have to stay 64-bit + std::array ple_layer_multipliers; + // head offsets and vocab sizes are token-space indices; the gather truncates them to int32 anyway + std::array ple_head_offsets; + std::array ple_head_vocab_sizes; + + bool is_ple(uint32_t il) const; + + // PLE conv history rows: (kernel - 1) * ngram_size; 0 without a PLE module + uint32_t ple_conv_state() const; + // qwen3vl deepstack // When parsed from GGUF, this implies the first N layers consume the first // N deepstack embeddings. Use deepstack_mapping_arr if you need a more diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index ec0f5a75314..8fafcd15304 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -6,12 +6,14 @@ #include "llama-context.h" #include +#include #include #include #include #include #include #include +#include static bool ggml_is_power_of_2(int n) { return (n & (n - 1)) == 0; @@ -77,7 +79,8 @@ llama_kv_cache::llama_kv_cache( llama_memory_t mem_other, const layer_filter_cb & filter, const layer_reuse_cb & reuse, - const layer_share_cb & share) : + const layer_share_cb & share, + const char * name_tag) : model(model), hparams(hparams), v_trans(v_trans), n_seq_max(n_seq_max), n_stream(unified ? 1 : n_seq_max), n_pad(n_pad), n_swa(n_swa), swa_type(swa_type), other(static_cast(mem_other)), @@ -231,8 +234,8 @@ llama_kv_cache::llama_kv_cache( ggml_tensor * k = has_k ? ggml_new_tensor_3d(ctx, type_k, n_embd_k_gqa, kv_size, n_stream) : nullptr; ggml_tensor * v = has_v ? ggml_new_tensor_3d(ctx, type_v, n_embd_v_gqa, kv_size, n_stream) : nullptr; - has_k && ggml_format_name(k, "cache_k_l%d", il); - has_v && ggml_format_name(v, "cache_v_l%d", il); + has_k && ggml_format_name(k, "cache_%sk_l%d", name_tag, il); + has_v && ggml_format_name(v, "cache_%sv_l%d", name_tag, il); std::vector k_stream; std::vector v_stream; @@ -1128,11 +1131,24 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & cells.pos_set(idx, ubatch.pos[i]); - if (ubatch.is_pos_2d()) { - llama_kv_cell_ext ext { - /*.x =*/ ubatch.pos[i + ubatch.n_tokens*2], - /*.y =*/ ubatch.pos[i + ubatch.n_tokens], - }; + if (ubatch.is_pos_2d() || ubatch.token || hparams.ple_n_heads > 0) { + llama_kv_cell_ext ext; + + if (ubatch.is_pos_2d()) { + ext.x = ubatch.pos[i + ubatch.n_tokens*2]; + ext.y = ubatch.pos[i + ubatch.n_tokens]; + } + + if (ubatch.token) { + ext.tok = ubatch.token[i]; + } else if (hparams.ple_n_heads > 0) { + // embd batch (multimodal input) has no token ids, need to pad it with the correct ID for PLE layers + // TODO @ngxson : check if we can do the same as gemma 3n / gemma 4 + ext.tok = hparams.ple_image_token_id != 0 + ? (llama_token) hparams.ple_image_token_id + : (llama_token) hparams.ple_eos_token_id; + } + cells.ext_set(idx, ext); } @@ -1805,6 +1821,115 @@ void llama_kv_cache::set_input_v_rot(ggml_tensor * dst) const { memcpy(dst->data, attn_rot_hadamard.at(n_rot).data(), ggml_nbytes(dst)); } +bool llama_kv_cache::has_cell_ext() const { + // M-RoPE needs the 2D position, the PLE n-gram hash needs the token id + return hparams.n_pos_per_embd() > 1 || hparams.ple_n_heads > 0; +} + +void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector & res) const { + const uint32_t n_tokens = ubatch.n_tokens; + + res.clear(); + res.resize(n_tokens*n, LLAMA_TOKEN_NULL); + + if (n == 0) { + return; + } + + // note: apply_ubatch() has already stored the current ubatch + // the window below thus covers tokens of this very ubatch as well, which is what we want + llama_pos p_min = std::numeric_limits::max(); + llama_pos p_max = std::numeric_limits::min(); + + std::bitset seqs; + + for (uint32_t i = 0; i < n_tokens; ++i) { + p_min = std::min(p_min, ubatch.pos[i]); + p_max = std::max(p_max, ubatch.pos[i]); + } + + for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) { + seqs.set(ubatch.seq_id_unq[s]); + } + + const llama_pos w0 = p_min - (llama_pos) n; + + // (seq_id, pos) -> token, for every cell that could be a predecessor of a ubatch token + std::unordered_map hist; + + const auto key = [](llama_seq_id seq_id, llama_pos pos) { + return ((uint64_t) seq_id << 32) | (uint32_t) pos; + }; + + // handle M-RoPE gaps: multiple tokens share the same temporal pos + // TODO @ngxson : improve this in the future + std::array, LLAMA_MAX_SEQ> below; + below.fill({ -1, LLAMA_TOKEN_NULL }); + + for (uint32_t s = 0; s < n_stream; ++s) { + // p_max inclusive: an embd token looks up cells at its own (shared) position + v_cells[s].for_each_token_in(seqs, 0, p_max + 1, + [&](llama_seq_id seq_id, llama_pos pos, llama_token tok) { + if (pos >= w0) { + hist[key(seq_id, pos)] = tok; + } else if (pos > below[seq_id].first) { + below[seq_id] = { pos, tok }; + } + }); + } + + // the token at pos p, or the nearest earlier one when p falls in an M-RoPE gap + const auto lookup = [&](llama_seq_id seq_id, llama_pos p) -> llama_token { + for (llama_pos q = p; q >= w0; --q) { + const auto it = hist.find(key(seq_id, q)); + if (it != hist.end()) { + return it->second; + } + } + return below[seq_id].second; + }; + + // an embd (multimodal) ubatch can repeat one position for a whole image, so positions + // do not encode the token order; resolve its predecessors by ubatch order instead + std::vector ord; // index among the ubatch tokens of the same seq + std::unordered_map> seq_idx; + + if (!ubatch.token) { + ord.resize(n_tokens); + for (uint32_t i = 0; i < n_tokens; ++i) { + auto & v = seq_idx[ubatch.seq_id[i][0]]; + ord[i] = v.size(); + v.push_back(i); + } + } + + for (uint32_t i = 0; i < n_tokens; ++i) { + // TODO: a token that belongs to more than one sequence has an ambiguous history. + // the n-gram architectures have to reject such batches + const llama_seq_id seq_id = ubatch.seq_id[i][0]; + + for (uint32_t j = 0; j < n; ++j) { + const llama_pos d = (llama_pos) (n - j); + + llama_pos p; + if (!ubatch.token) { + const auto & v = seq_idx[seq_id]; + const int64_t k = (int64_t) ord[i] - d; + // k >= 0: an earlier token of this very ubatch; k < 0: before the chunk + p = k >= 0 ? ubatch.pos[v[k]] : ubatch.pos[v[0]] + (llama_pos) k; + } else { + p = ubatch.pos[i] - d; + } + + if (p < 0) { + continue; + } + + res[i*n + j] = lookup(seq_id, p); + } + } +} + size_t llama_kv_cache::total_size() const { size_t size = 0; @@ -2037,6 +2162,15 @@ void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, lla } void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { + state_read_sinfo(io, seq_id, flags, nullptr, nullptr); +} + +void llama_kv_cache::state_read_sinfo( + llama_io_read_i & io, + llama_seq_id seq_id, + llama_state_seq_flags flags, + slot_info_vec_t * sinfos_out, +const slot_info_vec_t * sinfos_in) { // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] if (other) { return; @@ -2047,17 +2181,35 @@ void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama // TODO: fix incosistent handling of `seq_id < 0` and `seq_id == -1` in the codebase [TAG_LLAMA_SEQ_ID_NEG] GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size())); + if (sinfos_out) { + sinfos_out->assign(n_stream, slot_info{}); + } + + if (sinfos_in && sinfos_in->size() != n_stream) { + throw std::runtime_error("failed to restore kv cache: mirrored slot layout has the wrong stream count"); + } + uint32_t n_stream_cur; io.read(&n_stream_cur, sizeof(n_stream_cur)); if (n_stream_cur != n_stream) { throw std::runtime_error("n_stream mismatch"); } + // a whole-context restore replaces every stream, so the cache is emptied once here + // clear() resets all streams at once, so doing it per stream below would keep only the last one + if (seq_id == -1) { + clear(true); + } + for (uint32_t s = 0; s < n_stream; ++s) { uint32_t cell_count; io.read(&cell_count, sizeof(cell_count)); if (cell_count == 0) { + // a mirrored cache must be empty here as well, or the two no longer agree cell for cell + if (sinfos_in && !(*sinfos_in)[s].empty()) { + throw std::runtime_error("failed to restore kv cache: mirrored cache holds cells this one does not"); + } continue; } @@ -2066,7 +2218,7 @@ void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama slot_info sinfo; bool res = true; - res = res && state_read_meta(io, strm, cell_count, sinfo, seq_id); + res = res && state_read_meta(io, strm, cell_count, sinfo, seq_id, sinfos_in ? &(*sinfos_in)[s] : nullptr); try { res = res && state_read_data(io, strm, cell_count, sinfo); @@ -2082,6 +2234,10 @@ void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama } throw std::runtime_error("failed to restore kv cache"); } + + if (sinfos_out) { + (*sinfos_out)[s] = sinfo; + } } } @@ -2106,7 +2262,7 @@ void llama_kv_cache::state_write_meta(llama_io_write_i & io, const cell_ranges_t io.write(&pos, sizeof(pos)); io.write(&n_seq_id, sizeof(n_seq_id)); - if (hparams.n_pos_per_embd() > 1) { + if (has_cell_ext()) { const llama_kv_cell_ext ext = cells.ext_get(i); io.write(&ext, sizeof(ext)); } @@ -2217,7 +2373,7 @@ void llama_kv_cache::state_write_data(llama_io_write_i & io, const cell_ranges_t } } -bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, slot_info & sinfo, llama_seq_id dest_seq_id) { +bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, slot_info & sinfo, llama_seq_id dest_seq_id, const slot_info * sinfo_in) { auto & cells = v_cells[strm]; auto & head = v_heads[strm]; @@ -2243,12 +2399,17 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32 return false; } - if (hparams.n_pos_per_embd() > 1) { + if (has_cell_ext()) { llama_kv_cell_ext ext; io.read(&ext, sizeof(ext)); - ubatch.pos[i + ubatch.n_tokens] = ext.y; - ubatch.pos[i + ubatch.n_tokens*2] = ext.x; + if (hparams.n_pos_per_embd() > 1) { + ubatch.pos[i + ubatch.n_tokens] = ext.y; + ubatch.pos[i + ubatch.n_tokens*2] = ext.x; + } + + // apply_ubatch() below restores ext.tok from the ubatch tokens + ubatch.token[i] = ext.tok; } // read the sequence id, but directly discard it - we will use dest_seq_id instead @@ -2262,13 +2423,41 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32 ubatch.seq_id[i] = &dest_seq_id; } - sinfo = find_slot(ubatch, false); - if (sinfo.empty()) { - LLAMA_LOG_ERROR("%s: failed to find %d available cells in kv cache\n", __func__, cell_count); - return false; + if (sinfo_in) { + // this cache mirrors another one, so it takes that cache's layout instead of searching for its own cells + if (sinfo_in->empty() || sinfo_in->n_stream() != 1 || sinfo_in->idxs[0].size() != cell_count) { + LLAMA_LOG_ERROR("%s: mirrored slot layout holds %d cells, this cache restores %d\n", __func__, + sinfo_in->empty() ? 0 : (int) sinfo_in->idxs[0].size(), cell_count); + return false; + } + + sinfo = *sinfo_in; + + // the layout is cell indices, so it means the same in both caches only while their streams line up + sinfo.s0 = strm; + sinfo.s1 = strm; + sinfo.strm[0] = strm; + + // seq_rm above freed exactly the cells this sequence held + // anything else in the way is a cache that had already drifted, which this restore must not hide + for (uint32_t i = 0; i < cell_count; ++i) { + const uint32_t idx = sinfo.idxs[0][i]; + + if (idx >= cells.size() || !cells.is_empty(idx)) { + LLAMA_LOG_ERROR("%s: cell %u of the mirrored slot layout is not free\n", __func__, idx); + return false; + } + } + } else { + sinfo = find_slot(ubatch, false); + if (sinfo.empty()) { + LLAMA_LOG_ERROR("%s: failed to find %d available cells in kv cache\n", __func__, cell_count); + return false; + } } - // TODO: we cannot yet restore llama_kv_cell_ext as the apply_ubatch() does not support it yet + // note: apply_ubatch() rebuilds llama_kv_cell_ext from the ubatch + // only ext.tok and the M-RoPE 2D position round-trip through it // see: https://github.com/ggml-org/llama.cpp/pull/16825#issuecomment-3460868350 apply_ubatch(sinfo, ubatch); @@ -2290,7 +2479,12 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32 return false; } - clear(true); + // the cells go in from 0, so a mirrored cache lands on the same ones as long as it restores the same count. the layout itself carries no more information here + if (sinfo_in && (sinfo_in->empty() || sinfo_in->n_stream() != 1 || sinfo_in->idxs[0].size() != cell_count)) { + LLAMA_LOG_ERROR("%s: mirrored slot layout holds %d cells, this cache restores %d\n", __func__, + sinfo_in->empty() ? 0 : (int) sinfo_in->idxs[0].size(), cell_count); + return false; + } for (uint32_t i = 0; i < cell_count; ++i) { llama_pos pos; @@ -2301,7 +2495,7 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32 cells.pos_set(i, pos); - if (hparams.n_pos_per_embd() > 1) { + if (has_cell_ext()) { llama_kv_cell_ext ext; io.read(&ext, sizeof(ext)); cells.ext_set(i, ext); @@ -2652,3 +2846,7 @@ void llama_kv_cache_context::set_input_k_rot(ggml_tensor * dst) const { void llama_kv_cache_context::set_input_v_rot(ggml_tensor * dst) const { kv->set_input_v_rot(dst); } + +void llama_kv_cache_context::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector & res) const { + kv->get_prev_tokens(ubatch, n, res); +} diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 6cb6dbd2f98..c4d8699def1 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -112,7 +112,9 @@ class llama_kv_cache : public llama_memory_i { llama_memory_t mem_other, const layer_filter_cb & filter, const layer_reuse_cb & reuse, - const layer_share_cb & share); + const layer_share_cb & share, + // a model can hold more than one cache, so the tensor names have to stay unique + const char * name_tag = ""); ~llama_kv_cache() = default; @@ -166,6 +168,17 @@ class llama_kv_cache : public llama_memory_i { const llama_kv_cells & get_cells(llama_seq_id seq_id) const; + // state_read, plus the cells the restored tokens were placed in + // a cache that mirrors another one (the qwen4exp indexer) must not search for its own cells: two searches agree only by luck + // sinfos_out: if set, filled with the layout used; a stream with no cells leaves an empty entry + // sinfos_in : if set, the layout to use instead of searching. one entry per stream, cell count must match the blob + void state_read_sinfo( + llama_io_read_i & io, + llama_seq_id seq_id, + llama_state_seq_flags flags, + slot_info_vec_t * sinfos_out, + const slot_info_vec_t * sinfos_in); + // // graph_build API // @@ -219,6 +232,17 @@ class llama_kv_cache : public llama_memory_i { void set_input_k_rot(ggml_tensor * dst) const; void set_input_v_rot(ggml_tensor * dst) const; + // true if llama_kv_cell_ext holds information that has to survive a state save/restore + bool has_cell_ext() const; + + // for every token of the ubatch, the ids of the n tokens that precede it in its sequence + // example for M-RoPE image case: tokens A B X X X C, where X is a 3-token image at pos 2 spanning positions 2..4: + // tok: A B X X X C + // pos: 0 1 2 2 2 5 + // prev, n=2: A -> [NULL, NULL], B -> [NULL, A], 3rd X -> [X, X], C -> [X, X] + // note: used by n-gram input embeddings + void get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector & res) const; + private: const llama_model & model; const llama_hparams & hparams; @@ -318,7 +342,8 @@ class llama_kv_cache : public llama_memory_i { void state_write_meta(llama_io_write_i & io, const cell_ranges_t & cr, llama_seq_id seq_id = -1) const; void state_write_data(llama_io_write_i & io, const cell_ranges_t & cr) const; - bool state_read_meta(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, slot_info & sinfo, llama_seq_id dest_seq_id = -1); + // sinfo_in, when set, replaces the find_slot call: the cells are given by the caller + bool state_read_meta(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, slot_info & sinfo, llama_seq_id dest_seq_id = -1, const slot_info * sinfo_in = nullptr); bool state_read_data(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, const slot_info & sinfo); }; @@ -401,6 +426,9 @@ class llama_kv_cache_context : public llama_memory_context_i { void set_input_k_rot(ggml_tensor * dst) const; void set_input_v_rot(ggml_tensor * dst) const; + // see llama_kv_cache::get_prev_tokens() + void get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector & res) const; + private: llama_memory_status status; diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h index fddd31a0b21..5167c037db3 100644 --- a/src/llama-kv-cells.h +++ b/src/llama-kv-cells.h @@ -15,6 +15,10 @@ struct llama_kv_cell_ext { llama_pos x = 0; llama_pos y = 0; + // when tok = LLAMA_TOKEN_NULL when the cell is produced by embedding input (i.e. multimodal) + // use case: n-gram embeddings hash + llama_token tok = LLAMA_TOKEN_NULL; + // return true if the current 2D spatial position is greater than other bool is_2d_gt(llama_pos ox, llama_pos oy) const { return (y > oy) || (y == oy && x > ox); @@ -23,7 +27,7 @@ struct llama_kv_cell_ext { void reset() { static_assert(std::is_trivially_copyable_v); - memset(this, 0, sizeof(*this)); + *this = llama_kv_cell_ext{}; } }; @@ -305,6 +309,29 @@ class llama_kv_cells { return seq[i].test(seq_id); } + // gather the token ids of the cells in `seqs` with position in [p0, p1) + // the callback receives (seq_id, pos, token) for every such (cell, seq) pair + // note: used by n-gram input embeddings to recover the tokens preceding a ubatch + template + void for_each_token_in(const std::bitset & seqs, llama_pos p0, llama_pos p1, F && f) const { + for (const auto & i : used) { + if (pos[i] < p0 || pos[i] >= p1) { + continue; + } + + const auto m = seq[i] & seqs; + if (m.none()) { + continue; + } + + for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) { + if (m.test(s)) { + f(s, pos[i], ext[i].tok); + } + } + } + } + // note: call only if the cell is not empty and the seq_id is not in the cell void seq_add(uint32_t i, llama_seq_id seq_id) { assert(i < pos.size()); diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp new file mode 100644 index 00000000000..d4e59d77e57 --- /dev/null +++ b/src/llama-memory-hybrid-idx.cpp @@ -0,0 +1,465 @@ +#include "llama-memory-hybrid-idx.h" + +#include "llama-impl.h" +#include "llama-batch.h" +#include "llama-io.h" +#include "llama-model.h" + +#include +#include +#include +#include +#include + +// +// llama_memory_hybrid_idx +// + +llama_memory_hybrid_idx::llama_memory_hybrid_idx( + const llama_model & model, + /* attn */ + ggml_type type_k, + ggml_type type_v, + bool v_trans, + uint32_t kv_size, + uint32_t n_pad, + uint32_t n_swa, + llama_swa_type swa_type, + /* recurrent */ + ggml_type type_r, + ggml_type type_s, + uint32_t rs_size, + /* common */ + uint32_t n_seq_max, + uint32_t n_rs_seq, + bool offload, + bool unified, + /* layer filters */ + const layer_filter_cb & filter_attn, + const layer_filter_cb & filter_recr, + const layer_filter_cb & filter_idx) : + llama_memory_hybrid( + model, + type_k, type_v, v_trans, kv_size, n_pad, n_swa, swa_type, + type_r, type_s, rs_size, + n_seq_max, n_rs_seq, offload, unified, + filter_attn, filter_recr), + hparams_idx(model.hparams), + mem_idx(filter_idx == nullptr ? nullptr : [&] { + // MQA with a single key head of indexer_head_size, as llama_kv_cache_dsa shapes its own + std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), 1); + hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size; + + LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size); + + return new llama_kv_cache( + model, hparams_idx, type_k, type_v, v_trans, offload, unified, + kv_size, n_seq_max, n_pad, n_swa, swa_type, + nullptr, filter_idx, nullptr, nullptr, "idx_"); + }()) {} + +llama_memory_context_ptr llama_memory_hybrid_idx::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) { + // note: repeats llama_memory_hybrid::init_batch, as the indexer needs the attention slot infos that the base context hides + do { + balloc.split_reset(); + + // follow the recurrent pattern for creating the ubatch splits + std::vector ubatches; + + while (true) { + llama_ubatch ubatch; + + if (embd_all) { + // if all tokens are output, split by sequence + ubatch = balloc.split_seq(n_ubatch); + } else { + // Use non-sequential split when KV cache is unified (needed for hellaswag/winogrande/multiple-choice) + const bool unified = (get_mem_attn()->get_n_stream() == 1); + + // [TAG_RECURRENT_ROLLBACK_SPLITS] + // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch + // so that the rollback snapshots remain valid + const uint32_t n_rs_seq = get_mem_recr()->n_rs_seq; + + ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0); + } + + if (ubatch.n_tokens == 0) { + break; + } + + ubatches.push_back(std::move(ubatch)); // NOLINT + } + + if (balloc.get_n_used() < balloc.get_n_tokens()) { + // failed to find a suitable split + break; + } + + // prepare the recurrent batches first + if (!get_mem_recr()->prepare(ubatches)) { + // TODO: will the recurrent cache be in an undefined context at this point? + LLAMA_LOG_ERROR("%s: failed to prepare recurrent ubatches\n", __func__); + return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); + } + + // prepare the attention cache + auto heads_attn = get_mem_attn()->prepare(ubatches); + if (heads_attn.empty()) { + LLAMA_LOG_ERROR("%s: failed to prepare attention ubatches\n", __func__); + return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); + } + + // the indexer uses the attention cache's slot layout; a separate one can drift from it + llama_kv_cache::slot_info_vec_t heads_idx; + if (mem_idx) { + heads_idx = heads_attn; + } + + return std::make_unique( + this, std::move(heads_attn), std::move(heads_idx), std::move(ubatches)); + } while(false); + + return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); +} + +llama_memory_context_ptr llama_memory_hybrid_idx::init_full() { + return std::make_unique(this); +} + +llama_memory_context_ptr llama_memory_hybrid_idx::init_update(llama_context * lctx, bool optimize) { + return std::make_unique(this, lctx, optimize); +} + +void llama_memory_hybrid_idx::clear(bool data) { + llama_memory_hybrid::clear(data); + + if (mem_idx) { + mem_idx->clear(data); + } +} + +bool llama_memory_hybrid_idx::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { + // same order as llama_memory_hybrid::seq_rm: the recurrent cache can refuse, so try it first + if (!get_mem_recr()->seq_rm(seq_id, p0, p1)) { + return false; + } + + if (mem_idx) { + mem_idx->seq_rm(seq_id, p0, p1); + } + + return get_mem_attn()->seq_rm(seq_id, p0, p1); +} + +void llama_memory_hybrid_idx::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { + llama_memory_hybrid::seq_cp(seq_id_src, seq_id_dst, p0, p1); + + if (mem_idx) { + mem_idx->seq_cp(seq_id_src, seq_id_dst, p0, p1); + } +} + +void llama_memory_hybrid_idx::seq_keep(llama_seq_id seq_id) { + llama_memory_hybrid::seq_keep(seq_id); + + if (mem_idx) { + mem_idx->seq_keep(seq_id); + } +} + +void llama_memory_hybrid_idx::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { + llama_memory_hybrid::seq_add(seq_id, p0, p1, shift); + + if (mem_idx) { + mem_idx->seq_add(seq_id, p0, p1, shift); + } +} + +void llama_memory_hybrid_idx::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { + llama_memory_hybrid::seq_div(seq_id, p0, p1, d); + + if (mem_idx) { + mem_idx->seq_div(seq_id, p0, p1, d); + } +} + +std::map llama_memory_hybrid_idx::memory_breakdown() const { + std::map mb = llama_memory_hybrid::memory_breakdown(); + + if (mem_idx) { + for (const auto & buft_size : mem_idx->memory_breakdown()) { + mb[buft_size.first] += buft_size.second; + } + } + + return mb; +} + +void llama_memory_hybrid_idx::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { + llama_memory_hybrid::state_write(io, seq_id, flags); + + // [TAG_HYBRID_IDX_STATE] the indexer section goes last, so it is a pure suffix: an old reader stops early instead of misparsing it + // The indexer mirrors the attention cache, so it uses the same PARTIAL_ONLY gate. + if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { + if (mem_idx) { + mem_idx->state_write(io, seq_id, flags); + } + } + +} + +void llama_memory_hybrid_idx::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { + // note: repeats llama_memory_hybrid::state_read + // the indexer needs the attention cache's cells, and a half-failed restore must leave all three caches alike + + // [TAG_HYBRID_IDX_SINFO] + // the indexer restore adopts the attention cache's layout instead of searching for cells of its own + // two find_slot calls agree only while both caches see the same occupancy, which a restore cannot promise + llama_kv_cache::slot_info_vec_t sinfos_attn; + + try { + if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { + get_mem_attn()->state_read_sinfo(io, seq_id, flags, mem_idx ? &sinfos_attn : nullptr, nullptr); + } + + get_mem_recr()->state_read(io, seq_id, flags); + + // [TAG_HYBRID_IDX_STATE] must mirror the write order in state_write + if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { + if (mem_idx) { + mem_idx->state_read_sinfo(io, seq_id, flags, nullptr, &sinfos_attn); + } + } + + } catch (...) { + // a half-restored context is the one state the indexer cannot fix by itself: attention holds new cells, the indexer old ones + // drop what was being restored from all of them, which is a state they do agree on. + state_drop(seq_id); + + throw; + } +} + +void llama_memory_hybrid_idx::state_drop(llama_seq_id seq_id) { + // dropped directly, not via seq_rm: the recurrent cache may refuse it and then only the other two get cleared + if (seq_id < 0) { + clear(true); + + return; + } + + get_mem_attn()->seq_rm(seq_id, -1, -1); + get_mem_recr()->seq_rm(seq_id, -1, -1); + + if (mem_idx) { + mem_idx->seq_rm(seq_id, -1, -1); + } +} + +llama_kv_cache * llama_memory_hybrid_idx::get_mem_idx() const { + return mem_idx.get(); +} + +// +// llama_memory_hybrid_idx_context +// + +// streams in each ubatch's slot info, matching get_k/get_v's `ns` +static std::vector llama_memory_hybrid_idx_ns(const llama_kv_cache::slot_info_vec_t & sinfos) { + std::vector res; + res.reserve(sinfos.size()); + + for (const auto & sinfo : sinfos) { + res.push_back(sinfo.s1 - sinfo.s0 + 1); + } + + return res; +} + +llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(llama_memory_status status) : + llama_memory_hybrid_context(status) {} + +llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(llama_memory_hybrid_idx * mem) : + llama_memory_hybrid_context(mem), + mem(mem), + // graph reservation walks a full context, and qwen4exp builds the sparse attention only when this is set + // without it the reserved worst case is the dense graph, so ggml-alloc must grow the buffer on the first decode + ns_ubatch(mem->get_mem_idx() == nullptr ? + std::vector() : std::vector{ mem->get_mem_idx()->get_n_stream() }), + ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : + new llama_kv_cache_context(mem->get_mem_idx())) {} + +llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context( + llama_memory_hybrid_idx * mem, + llama_context * lctx, + bool optimize) : + llama_memory_hybrid_context(mem, lctx, optimize), + mem(mem) {} + +llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context( + llama_memory_hybrid_idx * mem, + slot_info_vec_t sinfos_attn, + slot_info_vec_t sinfos_idx, + std::vector ubatches) : + // note: the base copies the ubatches; ctx_idx gets a copy of its own + llama_memory_hybrid_context(mem, std::move(sinfos_attn), ubatches), + mem(mem), + ns_ubatch(llama_memory_hybrid_idx_ns(sinfos_idx)), + ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : + new llama_kv_cache_context(mem->get_mem_idx(), std::move(sinfos_idx), ubatches)) {} + +bool llama_memory_hybrid_idx_context::next() { + if (ctx_idx) { + ctx_idx->next(); + } + + ++i_cur; + + return llama_memory_hybrid_context::next(); +} + +bool llama_memory_hybrid_idx_context::apply() { + bool res = llama_memory_hybrid_context::apply(); + + if (ctx_idx) { + res = res & ctx_idx->apply(); + } + + return res; +} + +const llama_kv_cache_context * llama_memory_hybrid_idx_context::get_idx() const { + return static_cast(ctx_idx.get()); +} + +uint32_t llama_memory_hybrid_idx_context::get_n_stream() const { + GGML_ASSERT(i_cur < ns_ubatch.size()); + + return ns_ubatch[i_cur]; +} + +void llama_memory_hybrid_idx_context::set_input_qsa( + ggml_tensor * cell_blk, + ggml_tensor * blk_cells, + ggml_tensor * blk_pos, + ggml_tensor * bias, + const llama_ubatch * ubatch, + uint32_t ratio, + bool blk_bias) const { + GGML_ASSERT(ratio > 0); + GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr); + + GGML_ASSERT(ggml_backend_buffer_is_host(cell_blk->buffer)); + + const int64_t n_kv = cell_blk->ne[0]; + const int64_t n_ns = cell_blk->ne[1]; // streams in this ubatch + const int64_t n_blocks = blk_pos->ne[0]/(4*n_ns); + const int64_t n_tokens = ubatch->n_tokens; + const int64_t r = ratio; + + GGML_ASSERT(n_tokens % n_ns == 0); + const int64_t n_tps = n_tokens/n_ns; // tokens per stream + + int32_t * dst_cell_blk = (int32_t *) cell_blk->data; + int32_t * dst_blk_cells = (int32_t *) blk_cells->data; + int32_t * dst_blk_pos = (int32_t *) blk_pos->data; + float * dst_bias = (float *) bias->data; + + // block b covers [b*ratio, (b+1)*ratio), so its first token is at b*ratio + // all mrope sections carry it: exact for text, approximate for images + for (int64_t sec = 0; sec < 4; ++sec) { + for (int64_t s = 0; s < n_ns; ++s) { + for (int64_t b = 0; b < n_blocks; ++b) { + dst_blk_pos[sec*(n_blocks*n_ns) + s*n_blocks + b] = (int32_t) (b*r); + } + } + } + + // one pass per stream: cell j is a different token in each, so no mapping is shared + std::vector blk_of(n_kv); + std::vector filled(n_blocks); + + for (int64_t s = 0; s < n_ns; ++s) { + // ubatch index s*n_tps belongs to this stream; ask which cells array it uses + const llama_seq_id seq_of_stream = ubatch->seq_id[s*n_tps][0]; + const auto & cells = mem->get_mem_idx()->get_cells(seq_of_stream); + + int32_t * cur_cell_blk = dst_cell_blk + s*n_kv; + int32_t * cur_blk_cells = dst_blk_cells + s*(r*n_blocks); + + // an incomplete block cannot be pooled; the bias below forces those tail cells in + // -1 means no usable block, and block 0 only keeps the gather in range + std::fill(blk_of.begin(), blk_of.end(), -1); + std::fill(filled.begin(), filled.end(), 0); + std::fill(cur_blk_cells, cur_blk_cells + r*n_blocks, 0); + + // a cell no block covers needs its own -inf, which a per-block bias cannot carry + // every cache path keeps the position below the cell window, so this stays false + bool oor = false; + + for (int64_t j = 0; j < n_kv; ++j) { + if (cells.is_empty(j)) { + continue; + } + + const llama_pos p = cells.pos_get(j); + const int64_t b = p/r; + + if (b >= n_blocks) { + oor = true; + continue; + } + + blk_of[j] = (int32_t) b; + cur_blk_cells[b*r + (p%r)] = (int32_t) j; + filled[b]++; + } + + GGML_ASSERT((!blk_bias || !oor) && "qsa: cell position runs past the cell window"); + + // per-block mode keeps an unpooled cell's real block, so the block's own -inf reaches it + // per-cell mode carries that -inf itself and only needs the gather in range + for (int64_t j = 0; j < n_kv; ++j) { + if (blk_of[j] >= 0 && filled[blk_of[j]] < r && !blk_bias) { + blk_of[j] = -1; + } + cur_cell_blk[j] = blk_of[j] < 0 ? 0 : blk_of[j]; + } + + for (int64_t ii = 0; ii < n_tps; ++ii) { + const int64_t i = s*n_tps + ii; + const llama_seq_id seq_id = ubatch->seq_id[i][0]; + const llama_pos q = ubatch->pos[i]; + + // the tail is an incomplete block and is always visible, as in the reference + const llama_pos tail_start = (q + 1)/r*r; + + if (blk_bias) { + // a block sits wholly inside or outside the tail, so one value covers it + // the caller adds the attention mask, which drops empty, foreign and future cells + float * cur_blk_bias = dst_bias + i*n_blocks; + + for (int64_t b = 0; b < n_blocks; ++b) { + // finite, so it can never meet a -inf and produce a nan + cur_blk_bias[b] = b*r >= tail_start ? 1e9f : (filled[b] < r ? -INFINITY : 0.0f); + } + + continue; + } + + float * cur_bias = dst_bias + i*n_kv; + + for (int64_t j = 0; j < n_kv; ++j) { + float v = -INFINITY; + + if (!cells.is_empty(j) && cells.seq_has(j, seq_id) && cells.pos_get(j) <= q) { + // finite, so it can never meet a -inf and produce a nan + v = cells.pos_get(j) >= tail_start ? 1e9f : (blk_of[j] < 0 ? -INFINITY : 0.0f); + } + + cur_bias[j] = v; + } + } + } +} diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h new file mode 100644 index 00000000000..e3472646d0f --- /dev/null +++ b/src/llama-memory-hybrid-idx.h @@ -0,0 +1,156 @@ +#pragma once + +#include "llama-memory-hybrid.h" + +#include +#include + +// +// llama_memory_hybrid_idx +// + +// llama_memory_hybrid plus a third cache with one indexer key per token, for block-sparse attention (qwen4exp QSA) +// the indexer is a side buffer over the attention cells: same size, padding, streams and slots, so cell j is one token in both + +class llama_memory_hybrid_idx : public llama_memory_hybrid { +public: + llama_memory_hybrid_idx( + const llama_model & model, + /* attn */ + ggml_type type_k, + ggml_type type_v, + bool v_trans, + uint32_t kv_size, + uint32_t n_pad, + uint32_t n_swa, + llama_swa_type swa_type, + /* recurrent */ + ggml_type type_r, + ggml_type type_s, + uint32_t rs_size, + /* common */ + uint32_t n_seq_max, + uint32_t n_rs_seq, + bool offload, + bool unified, + /* layer filters */ + const layer_filter_cb & filter_attn, + const layer_filter_cb & filter_recr, + /* the indexer cache exists only if this is given */ + const layer_filter_cb & filter_idx); + + ~llama_memory_hybrid_idx() = default; + + // + // llama_memory_i + // + + llama_memory_context_ptr init_batch( + llama_batch_allocr & balloc, + uint32_t n_ubatch, + bool embd_all) override; + + llama_memory_context_ptr init_full() override; + + llama_memory_context_ptr init_update(llama_context * lctx, bool optimize) override; + + void clear(bool data) override; + + bool seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) override; + void seq_cp (llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) override; + void seq_keep(llama_seq_id seq_id) override; + void seq_add (llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) override; + void seq_div (llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) override; + + std::map memory_breakdown() const override; + + // state write/load + + void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override; + void state_read (llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override; + + // + // llama_memory_hybrid_idx specific API + // + + llama_kv_cache * get_mem_idx() const; // nullptr when the model carries no indexer + +private: + // forget seq_id (all of it if seq_id < 0) in every cache at once, so a failed restore cannot leave the caches out of step + // seq_id < 0 drops the whole context, as the caches themselves do on a failed restore + void state_drop(llama_seq_id seq_id); + + // the indexer cache holds one key head per layer, so it needs its own hparams: + // llama_kv_cache keeps a reference to what it is given + llama_hparams hparams_idx; + + const std::unique_ptr mem_idx; +}; + +class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { +public: + using slot_info_vec_t = llama_kv_cache::slot_info_vec_t; + + // used for errors + explicit llama_memory_hybrid_idx_context(llama_memory_status status); + + // used to create a full-cache context + explicit llama_memory_hybrid_idx_context(llama_memory_hybrid_idx * mem); + + // used to create an update context + llama_memory_hybrid_idx_context( + llama_memory_hybrid_idx * mem, + llama_context * lctx, + bool optimize); + + // used to create a batch processing context from a batch + llama_memory_hybrid_idx_context( + llama_memory_hybrid_idx * mem, + slot_info_vec_t sinfos_attn, + slot_info_vec_t sinfos_idx, + std::vector ubatches); + + ~llama_memory_hybrid_idx_context() = default; + + // + // llama_memory_context_i + // + + bool next() override; + bool apply() override; + + // + // llama_memory_hybrid_idx_context specific API + // + + // nullptr with no indexer, and for the update context, which builds no sparse graph + const llama_kv_cache_context * get_idx() const; + + // streams in the current slot info, the `ns` of get_k/get_v; 1 if unified + uint32_t get_n_stream() const; + + // block-compressed sparse attention (qwen4exp QSA) over the cells of the indexer cache. + // Blocks cut the position line, not the cell array, so no caller assumes a contiguous layout: + // cell_blk I32 [n_kv, ns] block each cell belongs to + // blk_cells I32 [ratio*n_blocks, ns] cells making up each block + // blk_pos I32 [4*n_blocks*ns] mrope position rows of each block's first token + // bias F32 [n_kv, n_tokens/ns, ns] -inf where invisible, large where always visible + // blk_bias asks for the bias per block instead: [n_blocks, n_tokens/ns, ns] + // the caller then adds the attention mask, the only part of the bias that varies within a block + void set_input_qsa(ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos, + ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio, + bool blk_bias) const; + +private: + const llama_memory_hybrid_idx * mem = nullptr; + + // streams per ubatch, read from the slot infos before ctx_idx takes them + // declared first, so it is initialised while sinfos_idx is still intact + const std::vector ns_ubatch; + + // null unless the model has an indexer and this is a batch or full context + const llama_memory_context_ptr ctx_idx; + + // mirrors the base class's ubatch cursor, which is private there + size_t i_cur = 0; +}; diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index e2990972ef7..57919accf09 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -51,7 +51,8 @@ llama_memory_recurrent::llama_memory_recurrent( auto it = ctx_map.find(buft); if (it == ctx_map.end()) { ggml_init_params params = { - /*.mem_size =*/ size_t(2u*n_layer*ggml_tensor_overhead()), + // r and s per layer, plus the separate PLE conv row where the model has one + /*.mem_size =*/ size_t((hparams.ple_conv_state() > 0 ? 3u : 2u)*n_layer*ggml_tensor_overhead()), /*.mem_buffer =*/ NULL, /*.no_alloc =*/ true, }; @@ -71,6 +72,7 @@ llama_memory_recurrent::llama_memory_recurrent( r_l.resize(n_layer); s_l.resize(n_layer); + p_l.resize(n_layer); for (int i = 0; i < n_layer; i++) { if (filter && !filter(i)) { @@ -103,6 +105,13 @@ llama_memory_recurrent::llama_memory_recurrent( ggml_format_name(s, "cache_s_l%d", i); r_l[i] = r; s_l[i] = s; + + // the PLE history needs its own row: Meta must mirror it while the delta-net conv state next door stays split + if (hparams.ple_conv_state() > 0 && hparams.is_ple(i)) { + ggml_tensor * p = ggml_new_tensor_2d(ctx, type_r, hparams.ple_conv_state(), n_rows); + ggml_format_name(p, "cache_ple_r_l%d", i); + p_l[i] = p; + } } // allocate tensors and initialize the buffers to avoid NaNs in the padding @@ -119,11 +128,13 @@ llama_memory_recurrent::llama_memory_recurrent( { const size_t memory_size_r = size_r_bytes(); const size_t memory_size_s = size_s_bytes(); + const size_t memory_size_p = size_p_bytes(); - LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u seqs %2u rs_seq), R (%s): %7.2f MiB, S (%s): %7.2f MiB\n", __func__, - (float)(memory_size_r + memory_size_s) / (1024.0f * 1024.0f), mem_size, n_layer, n_seq_max, n_rs_seq, + LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u seqs %2u rs_seq), R (%s): %7.2f MiB, S (%s): %7.2f MiB, P (%s): %7.2f MiB\n", __func__, + (float)(memory_size_r + memory_size_s + memory_size_p) / (1024.0f * 1024.0f), mem_size, n_layer, n_seq_max, n_rs_seq, ggml_type_name(type_r), (float)memory_size_r / (1024.0f * 1024.0f), - ggml_type_name(type_s), (float)memory_size_s / (1024.0f * 1024.0f)); + ggml_type_name(type_s), (float)memory_size_s / (1024.0f * 1024.0f), + ggml_type_name(type_r), (float)memory_size_p / (1024.0f * 1024.0f)); } } @@ -740,6 +751,18 @@ size_t llama_memory_recurrent::size_s_bytes() const { return size_s_bytes; } +size_t llama_memory_recurrent::size_p_bytes() const { + size_t size_p_bytes = 0; + + for (const auto & p : p_l) { + if (p != nullptr) { + size_p_bytes += ggml_nbytes(p); + } + } + + return size_p_bytes; +} + void llama_memory_recurrent::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { GGML_UNUSED(flags); @@ -899,6 +922,17 @@ void llama_memory_recurrent::state_write_data(llama_io_write_i & io, const std:: const size_t buf_size = range_size * r_size_row; io.write_tensor(r_l[il], range.first * r_size_row, buf_size); } + + // the PLE conv history is a second recurrent row, so it has to travel with the first + if (p_l[il] != nullptr) { + const uint64_t p_size_row = ggml_row_size(p_l[il]->type, hparams.ple_conv_state()); + io.write(&p_size_row, sizeof(p_size_row)); + + for (const auto & range : cell_ranges) { + const size_t range_size = range.second - range.first; + io.write_tensor(p_l[il], range.first * p_size_row, range_size * p_size_row); + } + } } if (!s_trans) { @@ -1097,6 +1131,20 @@ bool llama_memory_recurrent::state_read_data(llama_io_read_i & io, uint32_t cell // Read and set the keys for the whole cell range io.read_tensor(r_l[il], head * r_size_row, cell_count * r_size_row); } + + if (p_l[il] != nullptr) { + uint64_t p_size_row_ref; + io.read(&p_size_row_ref, sizeof(p_size_row_ref)); + const size_t p_size_row = ggml_row_size(p_l[il]->type, hparams.ple_conv_state()); + if (p_size_row != p_size_row_ref) { + LLAMA_LOG_ERROR("%s: mismatched ple row size (%zu != %zu, layer %d)\n", __func__, p_size_row, (size_t) p_size_row_ref, il); + return false; + } + + if (cell_count) { + io.read_tensor(p_l[il], head * p_size_row, cell_count * p_size_row); + } + } } if (!s_trans) { @@ -1251,6 +1299,10 @@ ggml_tensor * llama_memory_recurrent_context::get_s_l(int32_t il) const { return mem->s_l[il]; } +ggml_tensor * llama_memory_recurrent_context::get_p_l(int32_t il) const { + return mem->p_l[il]; +} + int32_t llama_memory_recurrent_context::s_copy(int i) const { const uint32_t cell_idx = i + mem->head; const int32_t src0 = mem->cells[cell_idx].src0; diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index b13b7b748f5..4abb3f5cf5c 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -111,6 +111,8 @@ class llama_memory_recurrent : public llama_memory_i { // per layer std::vector r_l; std::vector s_l; + // a second conv history that must stay replicated across devices, so it cannot share the r row + std::vector p_l; private: //const llama_model & model; @@ -125,6 +127,7 @@ class llama_memory_recurrent : public llama_memory_i { size_t size_r_bytes() const; size_t size_s_bytes() const; + size_t size_p_bytes() const; void state_write_meta(llama_io_write_i & io, const std::vector> & cell_ranges, llama_seq_id seq_id = -1) const; void state_write_data(llama_io_write_i & io, const std::vector> & cell_ranges) const; @@ -170,6 +173,7 @@ class llama_memory_recurrent_context : public llama_memory_context_i { ggml_tensor * get_r_l(int32_t il) const; ggml_tensor * get_s_l(int32_t il) const; + ggml_tensor * get_p_l(int32_t il) const; int32_t s_copy(int i) const; diff --git a/src/llama-mmap.cpp b/src/llama-mmap.cpp index ed572da7fb5..4d183cbc9c4 100644 --- a/src/llama-mmap.cpp +++ b/src/llama-mmap.cpp @@ -438,11 +438,34 @@ void llama_file::write_u32(uint32_t val) const { pimpl->write_u32(val); } // llama_mmap +#if defined(_POSIX_MAPPED_FILES) || defined(_WIN32) +// merge `ranges` and return their complement within [0, limit) +static llama_mmap::ranges ranges_complement(llama_mmap::ranges ranges, size_t limit) { + llama_mmap::ranges res; + std::sort(ranges.begin(), ranges.end()); + + size_t pos = 0; + for (const auto & range : ranges) { + const size_t beg = std::min(range.first, limit); + const size_t end = std::min(range.second, limit); + if (beg > pos) { + res.emplace_back(pos, beg); + } + pos = std::max(pos, end); + } + if (pos < limit) { + res.emplace_back(pos, limit); + } + + return res; +} +#endif + struct llama_mmap::impl { #ifdef _POSIX_MAPPED_FILES std::vector> mapped_fragments; - impl(struct llama_file * file, size_t prefetch, bool numa) { + impl(struct llama_file * file, size_t prefetch, bool numa, const llama_mmap::ranges & lazy_ranges) { size = file->size(); int fd = file->file_id(); int flags = MAP_SHARED; @@ -452,19 +475,35 @@ struct llama_mmap::impl { LLAMA_LOG_WARN("warning: posix_fadvise(.., POSIX_FADV_SEQUENTIAL) failed: %s\n", strerror(errno)); } - if (prefetch) { flags |= MAP_POPULATE; } + // MAP_POPULATE would fault in the lazy ranges too + if (prefetch && lazy_ranges.empty()) { flags |= MAP_POPULATE; } #endif addr = mmap(NULL, file->size(), PROT_READ, flags, fd, 0); if (addr == MAP_FAILED) { throw std::runtime_error(format("mmap failed: %s", strerror(errno))); } + // page-aligned madvise over [beg, end), clamped to the file + auto advise = [&](size_t beg, size_t end, int advice, const char * name) { + const size_t page_size = sysconf(_SC_PAGESIZE); + beg = beg & ~(page_size - 1); + end = std::min((end + page_size - 1) & ~(page_size - 1), file->size()); + if (beg >= end) { + return; + } + if (posix_madvise((char *) addr + beg, end - beg, advice)) { + LLAMA_LOG_WARN("warning: posix_madvise(.., %s) failed: %s\n", name, strerror(errno)); + } + }; + if (prefetch > 0) { - if (posix_madvise(addr, std::min(file->size(), prefetch), POSIX_MADV_WILLNEED)) { - LLAMA_LOG_WARN("warning: posix_madvise(.., POSIX_MADV_WILLNEED) failed: %s\n", - strerror(errno)); + for (const auto & range : ranges_complement(lazy_ranges, std::min(file->size(), prefetch))) { + advise(range.first, range.second, POSIX_MADV_WILLNEED, "POSIX_MADV_WILLNEED"); } } + for (const auto & range : lazy_ranges) { + advise(range.first, range.second, POSIX_MADV_RANDOM, "POSIX_MADV_RANDOM"); + } if (numa) { if (posix_madvise(addr, file->size(), POSIX_MADV_RANDOM)) { LLAMA_LOG_WARN("warning: posix_madvise(.., POSIX_MADV_RANDOM) failed: %s\n", @@ -533,7 +572,7 @@ struct llama_mmap::impl { #elif defined(_WIN32) HANDLE hMapping = nullptr; - impl(struct llama_file * file, size_t prefetch, bool numa) { + impl(struct llama_file * file, size_t prefetch, bool numa, const llama_mmap::ranges & lazy_ranges) { GGML_UNUSED(numa); size = file->size(); @@ -563,10 +602,15 @@ struct llama_mmap::impl { pPrefetchVirtualMemory = (decltype(pPrefetchVirtualMemory))(void *) GetProcAddress(hKernel32, "PrefetchVirtualMemory"); if (pPrefetchVirtualMemory) { - WIN32_MEMORY_RANGE_ENTRY range; - range.VirtualAddress = addr; - range.NumberOfBytes = (SIZE_T) std::min(size, prefetch); - if (!pPrefetchVirtualMemory(GetCurrentProcess(), 1, &range, 0)) { + std::vector entries; + for (const auto & range : ranges_complement(lazy_ranges, std::min(size, prefetch))) { + WIN32_MEMORY_RANGE_ENTRY entry; + entry.VirtualAddress = (char *) addr + range.first; + entry.NumberOfBytes = (SIZE_T) (range.second - range.first); + entries.push_back(entry); + } + if (!entries.empty() && + !pPrefetchVirtualMemory(GetCurrentProcess(), (ULONG_PTR) entries.size(), entries.data(), 0)) { LLAMA_LOG_WARN("warning: PrefetchVirtualMemory failed: %s\n", llama_format_win_err(GetLastError()).c_str()); } @@ -597,10 +641,11 @@ struct llama_mmap::impl { } } #else - impl(struct llama_file * file, size_t prefetch, bool numa) { + impl(struct llama_file * file, size_t prefetch, bool numa, const llama_mmap::ranges & lazy_ranges) { GGML_UNUSED(file); GGML_UNUSED(prefetch); GGML_UNUSED(numa); + GGML_UNUSED(lazy_ranges); throw std::runtime_error("mmap not supported"); } @@ -617,7 +662,8 @@ struct llama_mmap::impl { size_t size; }; -llama_mmap::llama_mmap(struct llama_file * file, size_t prefetch, bool numa) : pimpl(std::make_unique(file, prefetch, numa)) {} +llama_mmap::llama_mmap(struct llama_file * file, size_t prefetch, bool numa, + const ranges & lazy_ranges) : pimpl(std::make_unique(file, prefetch, numa, lazy_ranges)) {} llama_mmap::~llama_mmap() = default; size_t llama_mmap::size() const { return pimpl->size; } diff --git a/src/llama-mmap.h b/src/llama-mmap.h index b7d5c61e95f..cc28c8a73fa 100644 --- a/src/llama-mmap.h +++ b/src/llama-mmap.h @@ -2,6 +2,7 @@ #include #include +#include #include #include @@ -41,8 +42,12 @@ struct llama_file { }; struct llama_mmap { + // list of [first, last) byte ranges within a file + using ranges = std::vector>; + llama_mmap(const llama_mmap &) = delete; - llama_mmap(struct llama_file * file, size_t prefetch = (size_t) -1, bool numa = false); + llama_mmap(struct llama_file * file, size_t prefetch = (size_t) -1, bool numa = false, + const ranges & lazy_ranges = {}); ~llama_mmap(); size_t size() const; diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 9b22cb05f29..1b1f852a010 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -321,10 +321,11 @@ namespace GGUFMeta { case GGUF_TYPE_UINT32: case GGUF_TYPE_INT32: type_ok = (std::is_same::value) || (std::is_same::value); break; + case GGUF_TYPE_UINT64: type_ok = (std::is_same::value); break; case GGUF_TYPE_FLOAT32: type_ok = (std::is_same::value); break; case GGUF_TYPE_STRING: type_ok = (std::is_same::value); break; default: - throw std::runtime_error(format("%s is not a string/float32/uint32/int32 array", key.c_str())); + throw std::runtime_error(format("%s is not a string/float32/uint32/int32/uint64 array", key.c_str())); } if (!type_ok) { throw std::runtime_error(format("%s has wrong array element type %s", key.c_str(), gguf_type_name(arr_info.gt))); @@ -367,10 +368,11 @@ namespace GGUFMeta { case GGUF_TYPE_UINT32: case GGUF_TYPE_INT32: type_ok = (std::is_same::value) || (std::is_same::value); break; + case GGUF_TYPE_UINT64: type_ok = (std::is_same::value); break; case GGUF_TYPE_FLOAT32: type_ok = (std::is_same::value); break; case GGUF_TYPE_STRING: type_ok = (std::is_same::value); break; default: - throw std::runtime_error(format("%s is not a string/float32/uint32/int32 array", key.c_str())); + throw std::runtime_error(format("%s is not a string/float32/uint32/int32/uint64 array", key.c_str())); } if (!type_ok) { throw std::runtime_error(format("%s has wrong array element type %s", key.c_str(), gguf_type_name(arr_info.gt))); @@ -410,6 +412,9 @@ namespace GGUFMeta { template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::vector & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); + template bool llama_model_loader::get_arr>(enum llm_kv kid, std::vector & result, bool required); + template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); + template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); template bool llama_model_loader::get_key(const std::string & key, T & result, bool required) { @@ -1282,6 +1287,18 @@ struct ggml_tensor * llama_model_loader::create_tensor( return NULL; } + if ((flags & TENSOR_READ_LAZY) && use_mmap && lazy_mode != LLAMA_LAZY_MODE_OFF) { + // in auto mode, small tensors are cheap enough to keep resident + constexpr size_t auto_lazy_min_size = 4ull * 1024 * 1024 * 1024; + if (lazy_mode == LLAMA_LAZY_MODE_ON || ggml_nbytes(cur) > auto_lazy_min_size) { + const auto & w = require_weight(tn.str().c_str()); + lazy_tensor_ranges[w.idx].emplace_back(w.offs, w.offs + ggml_nbytes(cur)); + + LLAMA_LOG_INFO("%s: tensor %s (size = %zu MiB) lazy read enabled\n", + __func__, tn.str().c_str(), ggml_nbytes(cur)/1024/1024); + } + } + ggml_tensor t_meta = *cur; if (flags & TENSOR_ALLOW_RESHAPE) { for (size_t dim = 0; dim < GGML_MAX_DIMS; dim++) { @@ -1349,7 +1366,9 @@ void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps if (use_mmap) { mappings.reserve(files.size()); mmaps_used.reserve(files.size()); - for (const auto & file : files) { + for (uint32_t idx = 0; idx < files.size(); idx++) { + const auto & file = files[idx]; + bool is_numa = false; auto * dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU); @@ -1361,7 +1380,11 @@ void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps } } - std::unique_ptr mapping = std::make_unique(file.get(), prefetch ? -1 : 0, is_numa); + const auto it_lazy = lazy_tensor_ranges.find(idx); + static const llama_mmap::ranges no_lazy_ranges; + + std::unique_ptr mapping = std::make_unique(file.get(), prefetch ? -1 : 0, is_numa, + it_lazy != lazy_tensor_ranges.end() ? it_lazy->second : no_lazy_ranges); mmaps_used.emplace_back(mapping->size(), 0); if (mlock_mmaps) { std::unique_ptr mlock_mmap(new llama_mlock()); @@ -1400,27 +1423,26 @@ void llama_model_loader::unmap_weight(const llama_tensor_weight & w) const { mappings.at(w.idx)->unmap_fragment(w.offs, w.offs + ggml_nbytes(w.tensor)); } -void llama_model_loader::load_data_for(struct ggml_tensor * cur) const { - const auto & w = require_weight(ggml_get_name(cur)); +const void * llama_model_loader::load_data_range(const llama_tensor_weight & w, size_t offs, size_t size, void * buf) const { + GGML_ASSERT(offs + size <= ggml_nbytes(w.tensor)); + + const void * data = buf; if (use_mmap) { - const auto & mapping = mappings.at(w.idx); - if (cur->data == nullptr) { - cur->data = (uint8_t *)mapping->addr() + w.offs; - } else { - memcpy(cur->data, (uint8_t *)mapping->addr() + w.offs, ggml_nbytes(cur)); - } + data = (const uint8_t *) mappings.at(w.idx)->addr() + w.offs + offs; } else { - GGML_ASSERT(cur->data != nullptr); + GGML_ASSERT(buf != nullptr); GGML_ASSERT(w.idx < files.size()); const auto & file = files.at(w.idx); - file->seek(w.offs, SEEK_SET); - file->read_raw(cur->data, ggml_nbytes(cur)); + file->seek(w.offs + offs, SEEK_SET); + file->read_raw(buf, size); } - if (check_tensors && !ggml_validate_row_data(cur->type, cur->data, ggml_nbytes(cur))) { - throw std::runtime_error(format("tensor '%s' has invalid data", ggml_get_name(cur))); + if (check_tensors && !ggml_validate_row_data(w.tensor->type, data, size)) { + throw std::runtime_error(format("tensor '%s' has invalid data", ggml_get_name(w.tensor))); } + + return data; } bool llama_model_loader::load_all_data( diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index e9fe3592d42..20f74425389 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -68,6 +68,7 @@ struct llama_model_loader { static const int TENSOR_SKIP = 1 << 2; static const int TENSOR_SKIP_IF_VIRTUAL = 1 << 3; static const int TENSOR_ALLOW_RESHAPE = 1 << 4; + static const int TENSOR_READ_LAZY = 1 << 5; // read rows on demand instead of loading whole tensor; requires mmap for now int n_kv = 0; int n_tensors = 0; @@ -82,12 +83,18 @@ struct llama_model_loader { bool no_alloc; bool load_mtp; + // set by the caller before the create_tensor() calls + enum llama_lazy_mode lazy_mode = LLAMA_LAZY_MODE_OFF; + llama_files files; llama_ftype ftype; llama_fver fver; llama_mmaps mappings; + // byte ranges of TENSOR_READ_LAZY tensors, per file index + std::map lazy_tensor_ranges; + std::map weights_map; std::unordered_map kv_overrides; const llama_model_tensor_buft_override * tensor_buft_overrides; @@ -197,8 +204,9 @@ struct llama_model_loader { // release a weight's mmap pages void unmap_weight(const llama_tensor_weight & w) const; - // for backwards compatibility, does not support ggml-backend - void load_data_for(struct ggml_tensor * cur) const; + // read a byte range of a weight's data + // with mmap, returns a pointer into the mapping, otherwise reads into buf and returns buf + const void * load_data_range(const llama_tensor_weight & w, size_t offs, size_t size, void * buf) const; // Returns false if cancelled by progress_callback bool load_all_data( diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 9adaa93f62e..8860bd3f434 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -60,6 +60,10 @@ void llama_model_saver::add_kv(const enum llm_kv key, const int32_t value) { gguf_set_val_i32(gguf_ctx, llm_kv(key).c_str(), value); } +void llama_model_saver::add_kv(const enum llm_kv key, const uint64_t value) { + gguf_set_val_u64(gguf_ctx, llm_kv(key).c_str(), value); +} + void llama_model_saver::add_kv(const enum llm_kv key, const float value) { gguf_set_val_f32(gguf_ctx, llm_kv(key).c_str(), value); } @@ -113,6 +117,8 @@ void llama_model_saver::add_kv(const enum llm_kv key, const Container & value, c gguf_set_arr_data(gguf_ctx, llm_kv(key).c_str(), GGUF_TYPE_BOOL, value.data(), n_values); } else if (std::is_same::value) { gguf_set_arr_data(gguf_ctx, llm_kv(key).c_str(), GGUF_TYPE_INT32, value.data(), n_values); + } else if (std::is_same::value) { + gguf_set_arr_data(gguf_ctx, llm_kv(key).c_str(), GGUF_TYPE_UINT64, value.data(), n_values); } else if (std::is_same::value) { gguf_set_arr_data(gguf_ctx, llm_kv(key).c_str(), GGUF_TYPE_FLOAT32, value.data(), n_values); } else if (std::is_same::value) { @@ -124,6 +130,7 @@ void llama_model_saver::add_kv(const enum llm_kv key, const Container & value, c // instantiate for external usage: template void llama_model_saver::add_kv>(const enum llm_kv, const std::vector &, const bool); template void llama_model_saver::add_kv>(const enum llm_kv, const std::vector &, const bool); +template void llama_model_saver::add_kv>(const enum llm_kv, const std::vector &, const bool); void llama_model_saver::add_kv(const enum llm_kv key, const std::vector & value) { std::vector tmp(value.size()); @@ -308,6 +315,32 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters); add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps); add_kv(LLM_KV_HASH_LAYER_COUNT, hparams.dsv4_hash_layer_count); + add_kv(LLM_KV_HYPER_CONNECTION_LOW_RANK, hparams.hc_low_rank); + + // the PLE group only means anything whole: write all of it or none + if (hparams.ple_n_heads > 0) { + std::vector ple_layers; + for (uint32_t il = 0; il < hparams.n_layer_all; ++il) { + if (hparams.is_ple_impl[il]) { + ple_layers.push_back(il); + } + } + add_kv(LLM_KV_PLE_LAYERS, ple_layers); + add_kv(LLM_KV_PLE_NGRAM_SIZE, hparams.ple_ngram_size); + add_kv(LLM_KV_PLE_HEADS_PER_NGRAM, hparams.ple_heads_per_ngram); + add_kv(LLM_KV_PLE_CONV_KERNEL, hparams.ple_conv_kernel); + add_kv(LLM_KV_PLE_EOS_TOKEN_ID, hparams.ple_eos_token_id); + add_kv(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.ple_head_dim); + add_kv(LLM_KV_PLE_LAYER_MULTIPLIERS, std::vector( + hparams.ple_layer_multipliers.begin(), + hparams.ple_layer_multipliers.begin() + hparams.ple_ngram_size)); + add_kv(LLM_KV_PLE_HEAD_OFFSETS, std::vector( + hparams.ple_head_offsets.begin(), + hparams.ple_head_offsets.begin() + hparams.ple_n_heads)); + add_kv(LLM_KV_PLE_HEAD_VOCAB_SIZES, std::vector( + hparams.ple_head_vocab_sizes.begin(), + hparams.ple_head_vocab_sizes.begin() + hparams.ple_n_heads)); + } const float rope_scaling_factor = hparams.rope_freq_scale_train == 1.0f ? 0.0f : 1.0f/hparams.rope_freq_scale_train; @@ -442,6 +475,10 @@ void llama_model_saver::add_tensors_from_model() { add_tensor(model->hc_head_fn); add_tensor(model->hc_head_base); add_tensor(model->hc_head_scale); + add_tensor(model->per_layer_tok_embd); + add_tensor(model->hc_head_norm); + add_tensor(model->hc_head_down); + add_tensor(model->hc_head_up); for (const struct llama_layer & layer : model->layers) { for (size_t i = 0; i < sizeof(layer)/sizeof(struct ggml_tensor *); ++i) { diff --git a/src/llama-model-saver.h b/src/llama-model-saver.h index 36a715e2b6b..95e19e666e7 100644 --- a/src/llama-model-saver.h +++ b/src/llama-model-saver.h @@ -21,6 +21,7 @@ struct llama_model_saver { void add_kv(enum llm_kv key, uint32_t value); void add_kv(enum llm_kv key, int32_t value); + void add_kv(enum llm_kv key, uint64_t value); void add_kv(enum llm_kv key, float value); void add_kv(enum llm_kv key, bool value); void add_kv(enum llm_kv key, const char * value); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index c34700ff563..65a6702cefe 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -16,6 +16,7 @@ #include "llama-kv-cache-dsv4.h" #include "llama-memory-hybrid.h" #include "llama-memory-hybrid-iswa.h" +#include "llama-memory-hybrid-idx.h" #include "llama-memory-recurrent.h" #include "llama.h" @@ -319,6 +320,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_qwen35(params); case LLM_ARCH_QWEN35MOE: return new llama_model_qwen35moe(params); + case LLM_ARCH_QWEN4EXP: + return new llama_model_qwen4exp(params); case LLM_ARCH_MISTRAL3: return new llama_model_mistral3(params); case LLM_ARCH_EAGLE3: @@ -376,6 +379,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str static const std::regex pattern_qkv_bias ("blk\\.\\d*\\.attn_qkv.bias"); static const std::regex pattern_qk_norm ("blk\\.\\d*\\.attn_(q|k)_norm\\.weight"); static const std::regex pattern_kv_cache ("cache_(k|v)_l\\d*"); + static const std::regex pattern_idx_cache ("cache_idx_(k|v)_l\\d*"); static const std::regex pattern_dsv4_state ("dsv4_(csa|hca|lid)_state_(kv|score)_l\\d*"); static const std::regex pattern_attn_sinks ("blk\\.\\d*\\.attn_sinks.weight"); static const std::regex pattern_attn_out_weight ("blk\\.\\d*\\.attn_output.weight"); @@ -391,6 +395,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str static const std::regex pattern_ssm_beta ("blk\\.\\d*\\.ssm_beta.weight"); static const std::regex pattern_ssm_beta_alpha ("blk\\.\\d*\\.ssm_ba.weight"); static const std::regex pattern_r_cache ("cache_r_l\\d*"); + static const std::regex pattern_ple_r_cache ("cache_ple_r_l\\d*"); static const std::regex pattern_s_cache ("cache_s_l\\d*"); static const std::regex pattern_ssm_conv1d ("blk\\.\\d*\\.ssm_conv1d.weight"); static const std::regex pattern_ssm_out_weight ("blk\\.\\d*\\.ssm_out.weight"); @@ -488,6 +493,16 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str } } + // the qsa indexer has one key head and its projections are mirrored, so its cache cannot be split + if (std::regex_match(tensor_name, pattern_idx_cache)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); + } + + // the PLE table is model-level and its conv is mirrored, so every device runs the whole conv and needs the whole history + if (std::regex_match(tensor_name, pattern_ple_r_cache)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); + } + // standard attention if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_kv_weight)) { return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output.weight", "ssm_out.weight"); @@ -576,7 +591,8 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str }; auto get_split_segments = [&](int axis, uint32_t il) -> std::vector> { - if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE) { + if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE || + ud->model->arch == LLM_ARCH_QWEN4EXP) { const int64_t head_k_dim = hparams.ssm_d_state; const int64_t head_v_dim = hparams.ssm_d_state; const int64_t n_k_heads = hparams.ssm_n_group; @@ -714,7 +730,8 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_q_bias)) { GGML_ASSERT(segments.size() == 1); // some models have Q gate tensors, for those cases the granularity needs to be doubled: - if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE) { + if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE || + ud->model->arch == LLM_ARCH_QWEN4EXP) { return {std::lcm(2*n_embd_q, blck_size_perf)}; } return {granularity_q}; @@ -927,6 +944,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_35B_A3B: return "35B.A3B"; case LLM_TYPE_48B_A3B: return "48B.A3B"; case LLM_TYPE_80B_A3B: return "80B.A3B"; + case LLM_TYPE_A3B: return "A3B"; case LLM_TYPE_100B_A6B: return "100B.A6B"; case LLM_TYPE_102B_A12B: return "102B.A12B"; case LLM_TYPE_106B_A12B: return "106B.A12B"; @@ -2431,6 +2449,10 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, // layer filters, so pick the right one here llama_memory_hybrid::layer_filter_cb filter_attn = nullptr; llama_memory_hybrid::layer_filter_cb filter_recr = nullptr; + // only the sparse-attention architectures use llama_memory_hybrid_idx + // a null filter_idx means the GGUF has no indexer tensors + llama_memory_hybrid::layer_filter_cb filter_idx = nullptr; + const bool needs_mem_idx = (arch == LLM_ARCH_QWEN4EXP); if (arch == LLM_ARCH_FALCON_H1) { filter_attn = [&](uint32_t) { return true; }; filter_recr = [&](uint32_t) { return true; }; @@ -2441,13 +2463,20 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, filter_recr = [&](uint32_t il) { return hparams.is_recr(il) && hparams.n_ff(il) == 0; }; - } else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_MINIMAX_01) { + } else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_QWEN4EXP || arch == LLM_ARCH_MINIMAX_01) { filter_attn = [&](uint32_t il) { return il < hparams.n_layer() && !hparams.is_recr(il); }; filter_recr = [&](uint32_t il) { return il < hparams.n_layer() && hparams.is_recr(il); }; + + if (arch == LLM_ARCH_QWEN4EXP && hparams.indexer_head_size > 0) { + // QSA runs on the dense-attention layers only + filter_idx = [&](uint32_t il) { + return il < hparams.n_layer() && !hparams.is_recr(il); + }; + } } if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { @@ -2470,6 +2499,27 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, /* unified */ cparams.kv_unified, /* filter_attn */ std::move(filter_attn), /* filter_recr */ std::move(filter_recr)); + } else if (needs_mem_idx) { + // sparse attention over a per-token indexer cache, in its own memory type + res = new llama_memory_hybrid_idx( + /* model */ *this, + /* attn_type_k */ params.type_k, + /* attn_type_v */ params.type_v, + /* attn_v_trans */ !cparams.flash_attn, + /* attn_kv_size */ cparams.n_ctx_seq, + /* attn_n_pad */ 1, + /* attn_n_swa */ hparams.n_swa, + /* attn_swa_type */ hparams.swa_type, + /* recurrent_type_k */ GGML_TYPE_F32, + /* recurrent_type_v */ GGML_TYPE_F32, + /* recurrent_kv_size */ std::max((uint32_t) 1, cparams.n_seq_max), + /* n_seq_max */ cparams.n_seq_max, + /* n_rs_seq */ cparams.n_rs_seq, + /* offload */ cparams.offload_kqv, + /* unified */ cparams.kv_unified, + /* filter_attn */ std::move(filter_attn), + /* filter_recr */ std::move(filter_recr), + /* filter_idx */ std::move(filter_idx)); } else { res = new llama_memory_hybrid( /* model */ *this, @@ -2631,6 +2681,7 @@ llama_model_params llama_model_default_params() { /*.n_gpu_layers =*/ -1, /*.split_mode =*/ LLAMA_SPLIT_MODE_LAYER, /*.load_mode =*/ LLAMA_LOAD_MODE_AUTO, + /*.lazy_mode =*/ LLAMA_LAZY_MODE_AUTO, /*.main_gpu =*/ 0, /*.tensor_split =*/ nullptr, /*.progress_callback =*/ nullptr, @@ -2683,6 +2734,10 @@ int32_t llama_model_n_layer_nextn(const llama_model * model) { return model->hparams.n_layer_nextn; } +int32_t llama_model_dflash_selector_top_k(const llama_model * model) { + return model->hparams.dflash_selector_top_k; +} + int32_t llama_model_n_head(const llama_model * model) { return model->hparams.n_head(); } @@ -2881,6 +2936,10 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { return LLAMA_ROPE_TYPE_NEOX; case LLM_ARCH_DFLASH: + // drafts for M-RoPE targets carry rope sections and follow the target's temporal dim + if (const auto & s = model->hparams.rope_sections; s[0] || s[1] || s[2] || s[3]) { + return LLAMA_ROPE_TYPE_MROPE; + } // DSV4 DSpark drafters use DeepSeek-V4's normal RoPE; legacy DFlash backbones are NeoX return model->hparams.dsv4_hc_mult > 0 ? LLAMA_ROPE_TYPE_NORM : LLAMA_ROPE_TYPE_NEOX; @@ -2891,6 +2950,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_QWEN3VLMOE: case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: + case LLM_ARCH_QWEN4EXP: case LLM_ARCH_QWEN3TTS: return LLAMA_ROPE_TYPE_IMROPE; @@ -3067,7 +3127,8 @@ llama_model_base::llama_model_base(const struct llama_model_params & params) : l TENSOR_NOT_REQUIRED (llama_model_loader::TENSOR_NOT_REQUIRED), TENSOR_SKIP (llama_model_loader::TENSOR_SKIP), TENSOR_SKIP_IF_VIRTUAL(llama_model_loader::TENSOR_SKIP_IF_VIRTUAL), - TENSOR_ALLOW_RESHAPE (llama_model_loader::TENSOR_ALLOW_RESHAPE) {} + TENSOR_ALLOW_RESHAPE (llama_model_loader::TENSOR_ALLOW_RESHAPE), + TENSOR_READ_LAZY (llama_model_loader::TENSOR_READ_LAZY) {} ggml_tensor * llama_model_base::create_tensor(const LLM_TN_IMPL & tn, const std::initializer_list & ne, int flags) { GGML_ASSERT(ml != nullptr); diff --git a/src/llama-model.h b/src/llama-model.h index 44bd9675754..38066538ed1 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -129,6 +129,7 @@ enum llm_type { LLM_TYPE_35B_A3B, // Qwen3.5 LLM_TYPE_48B_A3B, // Kimi Linear LLM_TYPE_80B_A3B, // Qwen3 Next + LLM_TYPE_A3B, // Qwen3.8 Flash Next LLM_TYPE_100B_A6B, LLM_TYPE_102B_A12B, // Solar-Open LLM_TYPE_106B_A12B, // GLM-4.5-Air @@ -363,6 +364,11 @@ struct llama_layer { struct ggml_tensor * ffn_exp_probs_b = nullptr; struct ggml_tensor * ffn_gate_tid2eid = nullptr; + struct ggml_tensor * dflash_attn_conv_base = nullptr; + struct ggml_tensor * dflash_attn_conv_proj = nullptr; + struct ggml_tensor * dflash_ffn_conv_base = nullptr; + struct ggml_tensor * dflash_ffn_conv_proj = nullptr; + // mamba proj struct ggml_tensor * ssm_in = nullptr; struct ggml_tensor * ssm_x = nullptr; @@ -555,6 +561,22 @@ struct llama_layer { struct ggml_tensor * index_q_norm = nullptr; struct ggml_tensor * index_k_norm = nullptr; + struct ggml_tensor * hc_attn_norm = nullptr; + struct ggml_tensor * hc_attn_down = nullptr; + struct ggml_tensor * hc_attn_up = nullptr; + struct ggml_tensor * hc_attn_inject = nullptr; + struct ggml_tensor * hc_ffn_norm = nullptr; + struct ggml_tensor * hc_ffn_down = nullptr; + struct ggml_tensor * hc_ffn_up = nullptr; + struct ggml_tensor * hc_ffn_inject = nullptr; + + struct ggml_tensor * ple_key = nullptr; + struct ggml_tensor * ple_value = nullptr; + struct ggml_tensor * ple_norm_key = nullptr; + struct ggml_tensor * ple_norm_query = nullptr; + struct ggml_tensor * ple_norm_conv = nullptr; + struct ggml_tensor * ple_conv1d = nullptr; + // gemma4 layer output scale, reused for talkie embedding skip scale struct ggml_tensor * out_scale = nullptr; @@ -635,6 +657,10 @@ struct llama_model { struct ggml_tensor * altup_proj = nullptr; struct ggml_tensor * altup_unembd_proj = nullptr; struct ggml_tensor * per_layer_tok_embd = nullptr; + + struct ggml_tensor * hc_head_norm = nullptr; + struct ggml_tensor * hc_head_down = nullptr; + struct ggml_tensor * hc_head_up = nullptr; struct ggml_tensor * per_layer_model_proj = nullptr; struct ggml_tensor * per_layer_proj_norm = nullptr; @@ -646,9 +672,14 @@ struct llama_model { // dspark struct ggml_tensor * dspark_markov_w1 = nullptr; struct ggml_tensor * dspark_markov_w2 = nullptr; + struct ggml_tensor * dspark_markov_w2_s = nullptr; struct ggml_tensor * dspark_conf_proj = nullptr; struct ggml_tensor * dspark_conf_proj_b = nullptr; + struct ggml_tensor * dflash_selector_prev = nullptr; + struct ggml_tensor * dflash_selector_next = nullptr; + struct ggml_tensor * dflash_selector_hidden = nullptr; + // unified vector to store target-model extracted layer ids in eagle3, dflash, etc. std::vector target_layer_ids; @@ -756,6 +787,7 @@ struct llama_model_base : public llama_model { const int TENSOR_SKIP; const int TENSOR_SKIP_IF_VIRTUAL; const int TENSOR_ALLOW_RESHAPE; + const int TENSOR_READ_LAZY; explicit llama_model_base(const llama_model_params & params); virtual ~llama_model_base() = default; diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index 20252815d5c..c414caa173f 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -38,6 +38,9 @@ enum class tensor_category { OTHER }; +// max amount of tensor data kept in memory while quantizing a single tensor +static const size_t LLAMA_QUANT_MAX_BUF_SIZE = 8ull*1024*1024*1024; + static void zeros(std::ofstream & file, size_t n) { char zero = 0; for (size_t i = 0; i < n; ++i) { @@ -211,31 +214,26 @@ struct tensor_metadata { // static void llama_tensor_dequantize_impl( - ggml_tensor * tensor, std::vector> & output, std::vector & workers, + ggml_type type, const void * data, float * f32_output, std::vector & workers, const size_t nelements, const int nthread ) { - if (output.size() < nelements) { - output.resize(nelements); - } - float * f32_output = (float *) output.data(); - - const ggml_type_traits * qtype = ggml_get_type_traits(tensor->type); - if (ggml_is_quantized(tensor->type)) { + const ggml_type_traits * qtype = ggml_get_type_traits(type); + if (ggml_is_quantized(type)) { if (qtype->to_float == NULL) { - throw std::runtime_error(format("type %s unsupported for integer quantization: no dequantization available", ggml_type_name(tensor->type))); + throw std::runtime_error(format("type %s unsupported for integer quantization: no dequantization available", ggml_type_name(type))); } - } else if (tensor->type != GGML_TYPE_F16 && - tensor->type != GGML_TYPE_BF16) { - throw std::runtime_error(format("cannot dequantize/convert tensor type %s", ggml_type_name(tensor->type))); + } else if (type != GGML_TYPE_F16 && + type != GGML_TYPE_BF16) { + throw std::runtime_error(format("cannot dequantize/convert tensor type %s", ggml_type_name(type))); } if (nthread < 2) { - if (tensor->type == GGML_TYPE_F16) { - ggml_fp16_to_fp32_row((ggml_fp16_t *)tensor->data, f32_output, nelements); - } else if (tensor->type == GGML_TYPE_BF16) { - ggml_bf16_to_fp32_row((ggml_bf16_t *)tensor->data, f32_output, nelements); - } else if (ggml_is_quantized(tensor->type)) { - qtype->to_float(tensor->data, f32_output, nelements); + if (type == GGML_TYPE_F16) { + ggml_fp16_to_fp32_row((const ggml_fp16_t *)data, f32_output, nelements); + } else if (type == GGML_TYPE_BF16) { + ggml_bf16_to_fp32_row((const ggml_bf16_t *)data, f32_output, nelements); + } else if (ggml_is_quantized(type)) { + qtype->to_float(data, f32_output, nelements); } else { GGML_ABORT("fatal error"); // unreachable } @@ -243,14 +241,14 @@ static void llama_tensor_dequantize_impl( } size_t block_size; - if (tensor->type == GGML_TYPE_F16 || - tensor->type == GGML_TYPE_BF16) { + if (type == GGML_TYPE_F16 || + type == GGML_TYPE_BF16) { block_size = 1; } else { - block_size = (size_t)ggml_blck_size(tensor->type); + block_size = (size_t)ggml_blck_size(type); } - size_t block_size_bytes = ggml_type_size(tensor->type); + size_t block_size_bytes = ggml_type_size(type); GGML_ASSERT(nelements % block_size == 0); size_t nblocks = nelements / block_size; @@ -265,16 +263,16 @@ static void llama_tensor_dequantize_impl( size_t thr_elems = thr_blocks * block_size; // number of elements for this thread size_t thr_block_bytes = thr_blocks * block_size_bytes; // number of input bytes for this thread - auto compute = [qtype] (ggml_type typ, uint8_t * inbuf, float * outbuf, int nels) { + auto compute = [qtype] (ggml_type typ, const uint8_t * inbuf, float * outbuf, int nels) { if (typ == GGML_TYPE_F16) { - ggml_fp16_to_fp32_row((ggml_fp16_t *)inbuf, outbuf, nels); + ggml_fp16_to_fp32_row((const ggml_fp16_t *)inbuf, outbuf, nels); } else if (typ == GGML_TYPE_BF16) { - ggml_bf16_to_fp32_row((ggml_bf16_t *)inbuf, outbuf, nels); + ggml_bf16_to_fp32_row((const ggml_bf16_t *)inbuf, outbuf, nels); } else { qtype->to_float(inbuf, outbuf, nels); } }; - workers.emplace_back(compute, tensor->type, (uint8_t *) tensor->data + in_buff_offs, f32_output + out_buff_offs, thr_elems); + workers.emplace_back(compute, type, (const uint8_t *) data + in_buff_offs, f32_output + out_buff_offs, thr_elems); in_buff_offs += thr_block_bytes; out_buff_offs += thr_elems; } @@ -401,6 +399,12 @@ static ggml_type tensor_type_fallback(quantize_state_impl & qs, const ggml_tenso case GGML_TYPE_Q5_K: return_type = GGML_TYPE_Q5_1; break; case GGML_TYPE_Q6_K: return_type = GGML_TYPE_Q8_0; break; default: + if (qk_k <= 32) { + // the target is already a 32-block type, so there is no smaller block to demote to + // the check below turns it into F16, as a 256-block type does when its fallback does not fit + return_type = target_type; + break; + } throw std::runtime_error(format("no tensor type fallback is defined for type %s", ggml_type_name(target_type))); } @@ -681,7 +685,21 @@ static ggml_type llama_tensor_get_type(quantize_state_impl & qs, const llama_mod return tensor->type; } if (params->token_embedding_type < GGML_TYPE_COUNT && tm.category == tensor_category::TOKEN_EMBD) { - return params->token_embedding_type; + // per_layer_token_embd follows --token-embedding-type by default, but it is a large + // separate table, so let an explicit --tensor-type name it + bool named = false; + if (std::strcmp(tensor->name, "per_layer_token_embd.weight") == 0) { + const std::string tensor_name(tensor->name); + for (const auto & [pattern, qtype] : qs.tensor_type_patterns) { + if (std::regex_search(tensor_name, pattern)) { + named = true; + break; + } + } + } + if (!named) { + return params->token_embedding_type; + } } if (params->output_tensor_type < GGML_TYPE_COUNT && tm.category == tensor_category::OUTPUT) { return params->output_tensor_type; @@ -1093,6 +1111,8 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: std::vector> work; std::vector> f32_conv_buf; + const size_t max_buf_size = params->max_buf_size ? params->max_buf_size : LLAMA_QUANT_MAX_BUF_SIZE; + int cur_split = -1; std::ofstream fout; auto close_ofstream = [&]() { @@ -1143,15 +1163,13 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: const size_t tensor_size = ggml_nbytes(tensor); - if (!params->dry_run) { - if (!ml.use_mmap) { - if (read_data.size() < tensor_size) { - read_data.resize(tensor_size); - } - tensor->data = read_data.data(); + // read a byte range of the current tensor + auto load_range = [&](size_t offs, size_t size) -> const void * { + if (!ml.use_mmap && read_data.size() < size) { + read_data.resize(size); } - ml.load_data_for(tensor); - } + return ml.load_data_range(weight, offs, size, read_data.data()); + }; LLAMA_LOG_INFO("[%4d/%4d] %-36s - [%s], type = %6s, ", ++idx, ml.n_tensors, @@ -1166,7 +1184,6 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: // in then there's nothing to do. bool quantize = cur_type != new_type; - void * new_data; size_t new_size; if (params->dry_run) { @@ -1190,12 +1207,18 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: } else { // no --dry-run, perform quantization if (!quantize) { - new_data = tensor->data; new_size = tensor_size; LLAMA_LOG_INFO("size = %8.3f MiB\n", tensor_size/1024.0/1024.0); - } else { - const int64_t nelements = ggml_nelements(tensor); + // copy in slabs of whole rows, so that each slab can be validated + const size_t row_size = ggml_row_size(tensor->type, tensor->ne[0]); + const size_t slab_size = std::max(row_size, (max_buf_size/row_size)*row_size); + + for (size_t offs = 0; offs < tensor_size; offs += slab_size) { + const size_t size = std::min(slab_size, tensor_size - offs); + fout.write((const char *) load_range(offs, size), size); + } + } else { const float * imatrix = nullptr; if (imatrix_data) { auto it = imatrix_data->find(tm.remapped_imatrix_name); @@ -1227,43 +1250,60 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: throw std::runtime_error(format("Missing importance matrix for tensor %s in a very low-bit quantization", tensor->name)); } - float * f32_data; - - if (tensor->type == GGML_TYPE_F32) { - f32_data = (float *) tensor->data; - } else if (ggml_is_quantized(tensor->type) && !params->allow_requantize) { + if (ggml_is_quantized(tensor->type) && !params->allow_requantize) { throw std::runtime_error(format("requantizing from type %s is disabled", ggml_type_name(tensor->type))); - } else { - llama_tensor_dequantize_impl(tensor, f32_conv_buf, workers, nelements, nthread); - f32_data = (float *) f32_conv_buf.data(); } LLAMA_LOG_INFO("converting to %s .. ", ggml_type_name(new_type)); fflush(stdout); - if (work.size() < (size_t)nelements * 4) { - work.resize(nelements * 4); // upper bound on size - } - new_data = work.data(); - const int64_t n_per_row = tensor->ne[0]; const int64_t nrows = tensor->ne[1]; + const size_t row_size_src = ggml_row_size(tensor->type, n_per_row); + const size_t row_size_dst = ggml_row_size(new_type, n_per_row); + + // process the rows in slabs, so that the buffers stay below max_buf_size + const size_t bytes_per_row = row_size_src + row_size_dst + (tensor->type == GGML_TYPE_F32 ? 0 : n_per_row*sizeof(float)); + const int64_t nrows_slab = std::max(1, std::min(nrows, max_buf_size/bytes_per_row)); + static const int64_t min_chunk_size = 32 * 512; const int64_t chunk_size = (n_per_row >= min_chunk_size ? n_per_row : n_per_row * ((min_chunk_size + n_per_row - 1)/n_per_row)); - const int64_t nelements_matrix = tensor->ne[0] * tensor->ne[1]; - const int64_t nchunk = (nelements_matrix + chunk_size - 1)/chunk_size; - const int64_t nthread_use = nthread > 1 ? std::max((int64_t)1, std::min((int64_t)nthread, nchunk)) : 1; - // quantize each expert separately since they have different importance matrices new_size = 0; for (int64_t i03 = 0; i03 < tensor->ne[2]; ++i03) { - const float * f32_data_03 = f32_data + i03 * nelements_matrix; - void * new_data_03 = (char *)new_data + ggml_row_size(new_type, n_per_row) * i03 * nrows; const float * imatrix_03 = imatrix ? imatrix + i03 * n_per_row : nullptr; - new_size += llama_tensor_quantize_impl(new_type, f32_data_03, new_data_03, chunk_size, nrows, n_per_row, imatrix_03, workers, nthread_use); + for (int64_t ir = 0; ir < nrows; ir += nrows_slab) { + const int64_t nrows_cur = std::min(nrows_slab, nrows - ir); + const int64_t nelements_cur = nrows_cur * n_per_row; + + const void * src = load_range((i03*nrows + ir)*row_size_src, nrows_cur*row_size_src); + + const float * f32_data; + if (tensor->type == GGML_TYPE_F32) { + f32_data = (const float *) src; + } else { + if (f32_conv_buf.size() < (size_t) nelements_cur) { + f32_conv_buf.resize(nelements_cur); + } + llama_tensor_dequantize_impl(tensor->type, src, (float *) f32_conv_buf.data(), workers, nelements_cur, nthread); + f32_data = (const float *) f32_conv_buf.data(); + } + + if (work.size() < nrows_cur*row_size_dst) { + work.resize(nrows_cur*row_size_dst); + } + + const int64_t nchunk = (nelements_cur + chunk_size - 1)/chunk_size; + const int64_t nthread_use = nthread > 1 ? std::max((int64_t)1, std::min((int64_t)nthread, nchunk)) : 1; + + const size_t size_cur = llama_tensor_quantize_impl(new_type, f32_data, work.data(), chunk_size, nrows_cur, n_per_row, imatrix_03, workers, nthread_use); + + fout.write((const char *) work.data(), size_cur); + new_size += size_cur; + } } LLAMA_LOG_INFO("size = %8.2f MiB -> %8.2f MiB\n", tensor_size/1024.0/1024.0, new_size/1024.0/1024.0); } @@ -1273,10 +1313,8 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: // update the gguf metadata as we go gguf_set_tensor_type(ctx_outs[cur_split].get(), metadata[i].name.c_str(), new_type); GGML_ASSERT(gguf_get_tensor_size(ctx_outs[cur_split].get(), gguf_find_tensor(ctx_outs[cur_split].get(), metadata[i].name.c_str())) == new_size); - gguf_set_tensor_data(ctx_outs[cur_split].get(), metadata[i].name.c_str(), new_data); - // write tensor data + padding - fout.write((const char *) new_data, new_size); + // tensor data is already written, add the padding zeros(fout, GGML_PAD(new_size, align) - new_size); // unmap the tensor to free memory @@ -1323,7 +1361,8 @@ llama_model_quantize_params llama_model_quantize_default_params() { /*.imatrix =*/ nullptr, /*.kv_overrides =*/ nullptr, /*.tensor_type =*/ nullptr, - /*.prune_layers =*/ nullptr + /*.prune_layers =*/ nullptr, + /*.max_buf_size =*/ LLAMA_QUANT_MAX_BUF_SIZE }; return result; diff --git a/src/llama.cpp b/src/llama.cpp index 1609fec88dd..6ec5d315dce 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -318,6 +318,8 @@ static std::pair llama_model_load(struct gguf_context * meta llama_model_loader ml(metadata, set_tensor_data, set_tensor_data_ud, fname, splits, file, params.load_mode, params.check_tensors, params.no_alloc, params.load_mtp, params.kv_overrides, params.tensor_buft_overrides); + ml.lazy_mode = params.lazy_mode; + ml.print_info(); std::unique_ptr model_ptr(llama_model_create(ml, params)); diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index ff40c16b22e..f2c7d1d2462 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -7,6 +7,18 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale, false); + hparams.f_final_logit_softcapping = 0.0f; + ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false); + + // drafts for M-RoPE targets carry degenerate sections [n_rot/2, 0, 0, 0] + ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, false); + + ml.get_key(LLM_KV_DFLASH_BLOCK_SIZE, hparams.dflash_block_size, false); + ml.get_key(LLM_KV_DFLASH_CONV_KERNEL_SIZE, hparams.dflash_conv_kernel_size, false); + ml.get_key(LLM_KV_DFLASH_CONV_GROUP_SIZE, hparams.dflash_conv_group_size, false); + ml.get_key(LLM_KV_DFLASH_SELECTOR_RANK, hparams.dflash_selector_rank, false); + ml.get_key(LLM_KV_DFLASH_SELECTOR_TOP_K, hparams.dflash_selector_top_k, false); if (!ml.get_arr(LLM_KV_TARGET_LAYERS, target_layer_ids, false)) { throw std::runtime_error("DFlash model requires 'target_layers' in GGUF metadata"); @@ -103,15 +115,39 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { if (markov_meta) { const int64_t dspark_markov_rank = markov_meta->ne[0]; - dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { dspark_markov_rank, n_vocab }, 0); - dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab_draft }, 0); + dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { dspark_markov_rank, n_vocab }, 0); + dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab_draft }, 0); + dspark_markov_w2_s = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "scale"), { 1 }, TENSOR_NOT_REQUIRED); - dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, 0); + dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, TENSOR_NOT_REQUIRED); dspark_conf_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"), { 1 }, TENSOR_NOT_REQUIRED); LLAMA_LOG_INFO("%s: DFlash with DSpark markov head (rank = %lld)\n", __func__, (long long) dspark_markov_rank); } + const struct ggml_tensor * selector_meta = ml->get_tensor_meta("selector_hidden.weight"); + if (selector_meta) { + const int64_t rank = hparams.dflash_selector_rank; + if (rank <= 0 || hparams.dflash_block_size <= 0 || hparams.dflash_selector_top_k <= 0 || + hparams.dflash_conv_kernel_size <= 0 || hparams.dflash_conv_group_size <= 0) { + throw std::runtime_error("DFlash2 model is missing conv/selector metadata"); + } + if (n_embd % hparams.dflash_conv_group_size != 0) { + throw std::runtime_error("DFlash2 hidden size must be divisible by conv_group_size"); + } + if (n_embd < hparams.dflash_selector_top_k * (hparams.dflash_selector_top_k + 1)) { + throw std::runtime_error("DFlash2 hidden size is too small for the selector lattice"); + } + + dflash_selector_prev = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_PREV, "weight"), { rank, n_vocab }, 0); + dflash_selector_next = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_NEXT, "weight"), { rank, n_vocab }, 0); + dflash_selector_hidden = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, "weight"), { n_embd, rank }, 0); + + LLAMA_LOG_INFO("%s: DFlash2 conv kernel = %u, group = %u, selector rank = %u, top-k = %u\n", __func__, + hparams.dflash_conv_kernel_size, hparams.dflash_conv_group_size, + hparams.dflash_selector_rank, hparams.dflash_selector_top_k); + } + fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0); fc_s = create_tensor(tn(LLM_TENSOR_FC, "scale"), { 1 }, TENSOR_NOT_REQUIRED); output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc) @@ -184,10 +220,23 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, 0); layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), { n_embd_head_k }, 0); + // optional per-head attention sinks (e.g. Nemotron DSpark) + layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), { n_head }, TENSOR_NOT_REQUIRED); + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), { n_embd }, 0); layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), { n_embd, n_ff }, 0); layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, 0); layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, 0); + + if (selector_meta) { + const int64_t kernel = hparams.dflash_conv_kernel_size; + const int64_t groups = n_embd / hparams.dflash_conv_group_size; + const int64_t projected = 2 * kernel * groups; + layer.dflash_attn_conv_base = create_tensor(tn(LLM_TENSOR_DFLASH_ATTN_CONV_BASE, i), { n_embd, kernel, 2 }, 0); + layer.dflash_attn_conv_proj = create_tensor(tn(LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, "weight", i), { n_embd, projected }, 0); + layer.dflash_ffn_conv_base = create_tensor(tn(LLM_TENSOR_DFLASH_FFN_CONV_BASE, i), { n_embd, kernel, 2 }, 0); + layer.dflash_ffn_conv_proj = create_tensor(tn(LLM_TENSOR_DFLASH_FFN_CONV_PROJ, "weight", i), { n_embd, projected }, 0); + } } } @@ -245,7 +294,10 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & ggml_tensor * w1 = model.dspark_markov_w1; ggml_tensor * w2 = model.dspark_markov_w2; - GGML_ASSERT(w1 && w2 && model.dspark_conf_proj && "DSpark markov/confidence weights not loaded"); + GGML_ASSERT(w1 && w2 && "DSpark markov weights not loaded"); + + // confidence head is optional + const bool has_conf = model.dspark_conf_proj != nullptr; ggml_tensor * base = res->t_logits; // [n_vocab, n_tokens] const int64_t n_vocab = base->ne[0]; @@ -276,23 +328,22 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & ggml_tensor * prev = ggml_view_2d(ctx0, tokens, 1, n_blocks, token_stride, 0); prev = ggml_cont_1d(ctx0, prev, n_blocks); - // confidence head input: predicts per-position acceptance - ggml_tensor * conf_inp = res->t_embd; // [n_embd, n_tok] - ggml_tensor * cat = nullptr; ggml_tensor * cat_conf = nullptr; if (!sample_from_anchor) { // bonus anchor slot: pass the logits through unbiased, pad the (unread) confidence column - cat = ggml_cont(ctx0, ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, 0)); - cat_conf = ggml_sigmoid(ctx0, ggml_cont(ctx0, ggml_view_2d(ctx0, base, 1, n_blocks, base_stride, 0))); + cat = ggml_cont(ctx0, ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, 0)); + if (has_conf) { + cat_conf = ggml_sigmoid(ctx0, ggml_cont(ctx0, ggml_view_2d(ctx0, base, 1, n_blocks, base_stride, 0))); + } } // TODO: the in-graph chain is greedy (argmax); sampling params affect only the final // token pick, not the Markov conditioning path for (int64_t i = i_draft_beg; i < block_drafts; ++i) { - ggml_tensor * w1_prev = ggml_get_rows(ctx0, w1, prev); // [R, n_blocks] - ggml_tensor * bias = ggml_mul_mat(ctx0, w2, w1_prev); // [n_vocab_draft, n_blocks] + ggml_tensor * w1_prev = ggml_get_rows(ctx0, w1, prev); // [R, n_blocks] + ggml_tensor * bias = g.build_lora_mm(w2, w1_prev, model.dspark_markov_w2_s); // [n_vocab_draft, n_blocks] if (model.d2t) { // reduced draft vocab: scatter the bias to the target rows (base is -inf on the others) const int64_t n_draft_vocab = bias->ne[0]; @@ -309,17 +360,21 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & cat = cat ? ggml_concat(ctx0, cat, col, 1) : col; - // conf(i) = sigmoid(conf_proj . [conf_inp(i); markov_w1[prev(i)]] + b) -- [1, n_blocks] - ggml_tensor * conf_inp_i = ggml_view_2d(ctx0, conf_inp, conf_inp->ne[0], n_blocks, - (size_t) block_drafts * conf_inp->nb[1], i*conf_inp->nb[1]); - ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, conf_inp_i), w1_prev, 0); - ggml_tensor * conf = ggml_mul_mat(ctx0, model.dspark_conf_proj, feat); - if (model.dspark_conf_proj_b) { - conf = ggml_add(ctx0, conf, model.dspark_conf_proj_b); - } - conf = ggml_sigmoid(ctx0, conf); + if (has_conf) { + // confidence head input: predicts per-position acceptance + ggml_tensor * conf_inp = res->t_embd; // [n_embd, n_tok] + // conf(i) = sigmoid(conf_proj . [conf_inp(i); markov_w1[prev(i)]] + b) -- [1, n_blocks] + ggml_tensor * conf_inp_i = ggml_view_2d(ctx0, conf_inp, conf_inp->ne[0], n_blocks, + (size_t) block_drafts * conf_inp->nb[1], i*conf_inp->nb[1]); + ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, conf_inp_i), w1_prev, 0); + ggml_tensor * conf = ggml_mul_mat(ctx0, model.dspark_conf_proj, feat); + if (model.dspark_conf_proj_b) { + conf = ggml_add(ctx0, conf, model.dspark_conf_proj_b); + } + conf = ggml_sigmoid(ctx0, conf); - cat_conf = cat_conf ? ggml_concat(ctx0, cat_conf, conf, 1) : conf; + cat_conf = cat_conf ? ggml_concat(ctx0, cat_conf, conf, 1) : conf; + } if (i + 1 < block_drafts) { prev = ggml_argmax(ctx0, col); @@ -331,7 +386,7 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & out = ggml_cont(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3)); // [n_vocab, block_drafts, n_blocks] out = ggml_reshape_2d(ctx0, out, n_vocab, n_tok); - { + if (has_conf) { ggml_tensor * conf = ggml_reshape_3d(ctx0, cat_conf, 1, n_blocks, block_drafts); conf = ggml_cont(ctx0, ggml_permute(ctx0, conf, 0, 2, 1, 3)); conf = ggml_reshape_2d(ctx0, conf, 1, n_tok); @@ -346,6 +401,167 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & ggml_build_forward_expand(g.gf, out); } +static ggml_tensor * build_dflash2_conv( + llm_graph_context & g, + ggml_tensor * hidden, + ggml_tensor * dynamic, + ggml_tensor * base, + int side) { + const auto & hparams = g.hparams; + const int64_t hidden_size = hidden->ne[0]; + const int64_t n_tokens = hidden->ne[1]; + const int64_t n_blocks = g.ubatch.n_seqs_unq; + const int64_t kernel_size = hparams.dflash_conv_kernel_size; + const int64_t group_size = hparams.dflash_conv_group_size; + const int64_t n_groups = hidden_size / group_size; + + GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0); + GGML_ASSERT(dynamic && base && side >= 0 && side < 2); + + const int64_t block_size = n_tokens / n_blocks; + ggml_context * ctx0 = g.ctx0; + // ggml_cont copies even when the tensor is already contiguous + if (!ggml_is_contiguous(hidden) || hidden->ne[1] != n_tokens) { + hidden = ggml_cont_2d(ctx0, hidden, hidden_size, n_tokens); + } + if (!ggml_is_contiguous(dynamic) || dynamic->ne[1] != n_tokens) { + dynamic = ggml_cont_2d(ctx0, dynamic, dynamic->ne[0], n_tokens); + } + ggml_tensor * blocks = ggml_reshape_3d(ctx0, hidden, hidden_size, block_size, n_blocks); + ggml_tensor * coeffs = ggml_reshape_4d(ctx0, dynamic, n_groups, kernel_size, 2, n_tokens); + ggml_tensor * coeffs_side = ggml_view_3d(ctx0, coeffs, n_groups, kernel_size, n_tokens, + coeffs->nb[1], coeffs->nb[3], side * coeffs->nb[2]); + + ggml_tensor * coeff_all = ggml_cont(ctx0, coeffs_side); + coeff_all = ggml_reshape_4d(ctx0, coeff_all, 1, n_groups, kernel_size, n_tokens); + coeff_all = ggml_repeat_4d(ctx0, coeff_all, group_size, n_groups, kernel_size, n_tokens); + + ggml_tensor * base_side = ggml_reshape_4d(ctx0, + ggml_view_1d(ctx0, base, hidden_size * kernel_size, side * base->nb[2]), + group_size, n_groups, kernel_size, 1); + + ggml_tensor * weight_all = ggml_add(ctx0, coeff_all, base_side); + + ggml_tensor * result = nullptr; + for (int64_t tap = 0; tap < kernel_size; ++tap) { + ggml_tensor * values = blocks; + if (tap > 0) { + ggml_tensor * zeros = ggml_fill(ctx0, + ggml_new_tensor_3d(ctx0, hidden->type, hidden_size, std::min(tap, block_size), n_blocks), 0.0f); + if (tap < block_size) { + ggml_tensor * previous = ggml_view_3d(ctx0, blocks, hidden_size, block_size - tap, n_blocks, + blocks->nb[1], blocks->nb[2], 0); + values = ggml_concat(ctx0, zeros, previous, 1); + } else { + values = zeros; + } + } + values = ggml_reshape_2d(ctx0, values, hidden_size, n_tokens); + + ggml_tensor * weight = ggml_reshape_2d(ctx0, + ggml_cont(ctx0, ggml_view_4d(ctx0, weight_all, group_size, n_groups, 1, n_tokens, + weight_all->nb[1], weight_all->nb[2], weight_all->nb[3], tap * weight_all->nb[2])), + hidden_size, n_tokens); + + ggml_tensor * term = ggml_mul(ctx0, weight, values); + result = result ? ggml_add(ctx0, result, term) : term; + } + return result; +} + +// DFlash2 selector: top-k candidates per block position plus the pairwise +// transition scores, packed into the nextn output slot for the CPU-side walk. +static void build_dflash2_selector(llm_graph_context & g, const llama_model & model, ggml_tensor * tokens) { + ggml_context * ctx0 = g.ctx0; + auto & res = g.res; + + const auto & hparams = g.hparams; + const int64_t n_tokens = g.n_tokens; + const int64_t n_embd = g.n_embd; + + const int64_t top_k = hparams.dflash_selector_top_k; + const int64_t rank = hparams.dflash_selector_rank; + const int64_t n_blocks = g.ubatch.n_seqs_unq; + GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0); + GGML_ASSERT(res->t_logits->ne[1] == n_tokens); + if (!tokens) { + return; + } + + const int64_t tokens_per_block = n_tokens / n_blocks; + const int64_t block_size = std::min(tokens_per_block, hparams.dflash_block_size); + const int64_t row_used = top_k + top_k * top_k; + + ggml_tensor * candidates = ggml_top_k(ctx0, res->t_logits, top_k); + ggml_tensor * logits_rows = ggml_reshape_3d(ctx0, res->t_logits, 1, res->t_logits->ne[0], n_tokens); + ggml_tensor * unary = ggml_reshape_2d(ctx0, + ggml_get_rows(ctx0, logits_rows, candidates), top_k, n_tokens); + ggml_tensor * gate = g.build_lora_mm(model.dflash_selector_hidden, res->t_embd); + + // Everything below indexes [.., tokens_per_block, n_blocks]: the block + // position varies fastest, sequences are the outer dimension. + ggml_tensor * cand_blk = ggml_reshape_3d(ctx0, candidates, top_k, tokens_per_block, n_blocks); + ggml_tensor * unary_blk = ggml_reshape_3d(ctx0, unary, top_k, tokens_per_block, n_blocks); + ggml_tensor * gate_blk = ggml_reshape_3d(ctx0, gate, rank, tokens_per_block, n_blocks); + + // a position's score reads only the candidate sets at pos-1 and pos, so a run + // of positions has no internal dependency and scores in one batched matmul + auto score_run = [&](int64_t beg_pos, int64_t n_pos, ggml_tensor * pred_ids) { + ggml_tensor * cand_run = ggml_cont(ctx0, ggml_view_3d(ctx0, cand_blk, top_k, n_pos, n_blocks, + cand_blk->nb[1], cand_blk->nb[2], beg_pos * cand_blk->nb[1])); + ggml_tensor * unary_run = ggml_cont(ctx0, ggml_view_3d(ctx0, unary_blk, top_k, n_pos, n_blocks, + unary_blk->nb[1], unary_blk->nb[2], beg_pos * unary_blk->nb[1])); + ggml_tensor * gate_run = ggml_cont(ctx0, ggml_view_3d(ctx0, gate_blk, rank, n_pos, n_blocks, + gate_blk->nb[1], gate_blk->nb[2], beg_pos * gate_blk->nb[1])); + + const int64_t n_pred = pred_ids->ne[0] / (n_pos * n_blocks); + + ggml_tensor * successor = ggml_reshape_4d(ctx0, + ggml_get_rows(ctx0, model.dflash_selector_next, ggml_reshape_1d(ctx0, cand_run, top_k * n_pos * n_blocks)), + rank, top_k, n_pos, n_blocks); + ggml_tensor * predecessor = ggml_reshape_4d(ctx0, + ggml_get_rows(ctx0, model.dflash_selector_prev, pred_ids), + rank, n_pred, n_pos, n_blocks); + + ggml_tensor * gate_bcast = ggml_reshape_4d(ctx0, gate_run, rank, 1, n_pos, n_blocks); + ggml_tensor * cond = ggml_mul(ctx0, predecessor, ggml_repeat(ctx0, gate_bcast, predecessor)); + ggml_tensor * score = ggml_mul_mat(ctx0, successor, cond); + if (n_pred == 1) { + score = ggml_repeat_4d(ctx0, score, top_k, top_k, n_pos, n_blocks); + } + ggml_tensor * unary_bcast = ggml_reshape_4d(ctx0, unary_run, top_k, 1, n_pos, n_blocks); + score = ggml_add(ctx0, score, ggml_repeat(ctx0, unary_bcast, score)); + + ggml_tensor * row = ggml_concat(ctx0, + ggml_cast(ctx0, cand_run, GGML_TYPE_F32), + ggml_reshape_3d(ctx0, score, top_k * top_k, n_pos, n_blocks), 0); + return ggml_pad(ctx0, row, n_embd - row_used, 0, 0, 0); + }; + + ggml_tensor * packed = ggml_fill(ctx0, + ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_embd, 1, n_blocks), 0.0f); + + if (block_size > 1) { + // Position 1 alone: its predecessor is the anchor token, one id per + // sequence rather than a candidate set. + ggml_tensor * anchor_ids = ggml_cont_1d(ctx0, + ggml_view_2d(ctx0, tokens, 1, n_blocks, tokens_per_block * tokens->nb[0], 0), n_blocks); + packed = ggml_concat(ctx0, packed, score_run(1, 1, anchor_ids), 1); + } + if (block_size > 2) { + ggml_tensor * prev_ids = ggml_reshape_1d(ctx0, + ggml_cont(ctx0, ggml_view_3d(ctx0, cand_blk, top_k, block_size - 2, n_blocks, + cand_blk->nb[1], cand_blk->nb[2], cand_blk->nb[1])), + top_k * (block_size - 2) * n_blocks); + packed = ggml_concat(ctx0, packed, score_run(2, block_size - 2, prev_ids), 1); + } + + packed = ggml_reshape_2d(ctx0, packed, n_embd, block_size * n_blocks); + g.cb(packed, "dflash2_lattice", -1); + res->t_h_nextn = packed; + ggml_build_forward_expand(g.gf, packed); +} + // DFlash decoder, dual-mode by batch type: // * embd batch -> fused target features: project + inject K/V into the cache. // * token batch -> noise-block diffusion: attend over [committed, MASK...] to generate draft tokens @@ -370,6 +586,20 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra const float kq_scale = 1.0f/sqrtf(float(n_embd_head)); + // drafts for M-RoPE targets use degenerate sections (temporal dim only) + int sections[4]; + std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections); + + auto build_rope = [&](ggml_tensor * cur, ggml_tensor * pos) { + return rope_type == GGML_ROPE_TYPE_MROPE + ? ggml_rope_multi(ctx0, cur, pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow) + : ggml_rope_ext(ctx0, cur, pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + }; + // KV cache injection if (ubatch.embd) { auto inp = std::make_unique(n_embd); @@ -392,11 +622,7 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); - Kcur = ggml_rope_ext( - ctx0, Kcur, inp_pos, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow - ); + Kcur = build_rope(Kcur, inp_pos); cb(Kcur, "Kcur_injected", il); cb(Vcur, "Vcur_injected", il); @@ -450,6 +676,7 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); ggml_set_input(inp->tokens); + res->t_inp_tokens = inp->tokens; ggml_tensor * inp_tokens = inp->tokens; @@ -464,6 +691,13 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra ggml_tensor * noise_norm = build_norm(inpL, layer.attn_norm, NULL, LLM_NORM_RMS, il); cb(noise_norm, "noise_norm", il); + ggml_tensor * attn_dynamic = nullptr; + if (layer.dflash_attn_conv_proj) { + attn_dynamic = build_lora_mm(layer.dflash_attn_conv_proj, noise_norm); + noise_norm = build_dflash2_conv(*this, noise_norm, attn_dynamic, layer.dflash_attn_conv_base, 0); + cb(noise_norm, "attn_conv_in", il); + } + ggml_tensor * Qcur = build_lora_mm(layer.wq, noise_norm); ggml_tensor * Kcur = build_lora_mm(layer.wk, noise_norm); ggml_tensor * Vcur = build_lora_mm(layer.wv, noise_norm); @@ -475,24 +709,21 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra Qcur = build_norm(Qcur, layer.attn_q_norm, NULL, LLM_NORM_RMS, il); Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); - Qcur = ggml_rope_ext( - ctx0, Qcur, inp_pos, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow - ); - Kcur = ggml_rope_ext( - ctx0, Kcur, inp_pos, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow - ); + Qcur = build_rope(Qcur, inp_pos); + Kcur = build_rope(Kcur, inp_pos); cb(Qcur, "Qcur", il); cb(Kcur, "Kcur", il); cb(Vcur, "Vcur", il); // cache-aware, non-causal attention ggml_tensor * cur = use_iswa - ? build_attn(inp_attn_iswa, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il) - : build_attn(inp_attn, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + ? build_attn(inp_attn_iswa, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il) + : build_attn(inp_attn, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il); + + if (attn_dynamic) { + cur = build_dflash2_conv(*this, cur, attn_dynamic, layer.dflash_attn_conv_base, 1); + cb(cur, "attn_conv_out", il); + } ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL); cb(ffn_inp, "ffn_inp", il); @@ -500,6 +731,13 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il); cb(cur, "ffn_norm", il); + ggml_tensor * ffn_dynamic = nullptr; + if (layer.dflash_ffn_conv_proj) { + ffn_dynamic = build_lora_mm(layer.dflash_ffn_conv_proj, cur); + cur = build_dflash2_conv(*this, cur, ffn_dynamic, layer.dflash_ffn_conv_base, 0); + cb(cur, "ffn_conv_in", il); + } + cur = build_ffn(cur, layer.ffn_up, NULL, layer.ffn_up_s, layer.ffn_gate, NULL, layer.ffn_gate_s, @@ -508,6 +746,11 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra LLM_FFN_SILU, LLM_FFN_PAR, il); cb(cur, "ffn_out", il); + if (ffn_dynamic) { + cur = build_dflash2_conv(*this, cur, ffn_dynamic, layer.dflash_ffn_conv_base, 1); + cb(cur, "ffn_conv_out", il); + } + cur = ggml_add(ctx0, cur, ffn_inp); cb(cur, "l_out", il); @@ -532,6 +775,19 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra cur = build_lora_mm(output, cur, output_s); + // DFlash2 feeds these logits to the selector, so they need the target's output + // transforms; DFlash1 and DSpark read them through the sampler instead + if (model.dflash_selector_hidden) { + if (hparams.f_logit_scale != 0.0f) { + cur = ggml_scale(ctx0, cur, hparams.f_logit_scale); + } + if (hparams.f_final_logit_softcapping > 0.0f) { + cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping); + cur = ggml_tanh(ctx0, cur); + cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping); + } + } + // reduced-draft-vocab exports: scatter the draft logits to the target vocabulary via d2t if (model.d2t) { const int64_t n_draft_vocab = cur->ne[0]; @@ -556,6 +812,10 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra if (model.dspark_markov_w1) { build_dspark_markov_head(*this, model, inp_tokens); } + + if (model.dflash_selector_hidden) { + build_dflash2_selector(*this, model, inp_tokens); + } } // DSV4 DSpark decoder, dual-mode by batch type (see the DFlash decoder above): diff --git a/src/models/gemma4.cpp b/src/models/gemma4.cpp index e44f423bdbc..aa518c6df50 100644 --- a/src/models/gemma4.cpp +++ b/src/models/gemma4.cpp @@ -50,7 +50,7 @@ void llama_model_gemma4::load_arch_tensors(llama_model_loader &) { tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); if (n_embd_per_layer > 0) { - per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), {n_embd_per_layer * n_layer, n_vocab}, 0); + per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), {n_embd_per_layer * n_layer, n_vocab}, TENSOR_READ_LAZY); per_layer_model_proj = create_tensor(tn(LLM_TENSOR_PER_LAYER_MODEL_PROJ, "weight", 0), {n_embd, n_embd_per_layer * n_layer}, 0); per_layer_proj_norm = create_tensor(tn(LLM_TENSOR_PER_LAYER_PROJ_NORM, "weight", 0), {n_embd_per_layer}, 0); } diff --git a/src/models/minimax-01.cpp b/src/models/minimax-01.cpp index a6ccee1917e..361114acc32 100644 --- a/src/models/minimax-01.cpp +++ b/src/models/minimax-01.cpp @@ -174,16 +174,14 @@ class llm_graph_input_la : public llm_graph_input_i { bool can_reuse(const llm_graph_params & params) override { bool res = true; - if (params.ubatch.n_seq_tokens > 1) { - res &= ( inp_q_decay && inp_q_decay->ne[2] == params.ubatch.n_seq_tokens); - res &= ( inp_k_decay && inp_k_decay->ne[2] == params.ubatch.n_seq_tokens); - res &= (inp_diag_decay && inp_diag_decay->ne[1] == params.ubatch.n_seq_tokens); - } + res &= ( inp_q_decay && inp_q_decay->ne[2] == params.ubatch.n_seq_tokens); + res &= ( inp_k_decay && inp_k_decay->ne[2] == params.ubatch.n_seq_tokens); + res &= (inp_diag_decay && inp_diag_decay->ne[1] == params.ubatch.n_seq_tokens); return res; } - const llama_hparams & hparams; + const llama_hparams hparams; ggml_tensor * inp_slopes = nullptr; // F32 [n_head] ggml_tensor * inp_q_decay = nullptr; // F32 [1, n_head, n_batch] @@ -223,19 +221,17 @@ llama_model_minimax_01::graph::graph(const llama_model & model, const llm_graph_ ggml_set_input(inp->inp_slopes); cb(inp->inp_slopes, "slopes", -1); - if (n_seq_tokens != 1) { - inp->inp_q_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs); - ggml_set_input(inp->inp_q_decay); - cb(inp->inp_q_decay, "q_decay_exp", -1); + inp->inp_q_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs); + ggml_set_input(inp->inp_q_decay); + cb(inp->inp_q_decay, "q_decay_exp", -1); - inp->inp_k_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs); - ggml_set_input(inp->inp_k_decay); - cb(inp->inp_k_decay, "k_decay_exp", -1); + inp->inp_k_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs); + ggml_set_input(inp->inp_k_decay); + cb(inp->inp_k_decay, "k_decay_exp", -1); - inp->inp_diag_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_seq_tokens, n_seq_tokens, n_head, n_seqs); - ggml_set_input(inp->inp_diag_decay); - cb(inp->inp_diag_decay, "diag_decay_exp", -1); - } + inp->inp_diag_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_seq_tokens, n_seq_tokens, n_head, n_seqs); + ggml_set_input(inp->inp_diag_decay); + cb(inp->inp_diag_decay, "diag_decay_exp", -1); la = (llm_graph_input_la *) res->add_input(std::move(inp)); @@ -319,41 +315,8 @@ llama_model_minimax_01::graph::graph(const llama_model & model, const llm_graph_ ggml_tensor * qkv = nullptr; ggml_tensor * kv_new = nullptr; - - if (n_seq_tokens == 1) { - // lightning attention - optimized single token case for TG - - ggml_tensor * slopes_neg = ggml_scale(ctx0, slope_rate, -1.0); - cb(slopes_neg, "slopes_neg", il); - - ggml_tensor * ratio = ggml_exp(ctx0, slopes_neg); - cb(ratio, "ratio", il); - - ggml_tensor * ratio_3d = ggml_reshape_3d(ctx0, ratio, 1, 1, n_head); - cb(ratio_3d, "ratio3d", il); - - ggml_tensor * v_trans = ggml_cont(ctx0, ggml_permute(ctx0, Vcur, 1, 2, 0, 3)); - cb(v_trans, "v_trans", il); - - ggml_tensor * k_trans = ggml_cont(ctx0, ggml_permute(ctx0, Kcur, 1, 2, 0, 3)); - cb(k_trans, "k_trans", il); - - ggml_tensor * kv_cur = ggml_mul_mat(ctx0, k_trans, v_trans); - cb(kv_cur, "kv_cur", il); - - ggml_tensor * kv_old_s = ggml_mul(ctx0, kv_old, ratio_3d); - cb(kv_old_s, "kv_old_s", il); - - kv_new = ggml_add(ctx0, kv_old_s, kv_cur); - cb(kv_new, "kv_new", il); - - ggml_tensor * q_trans = ggml_permute(ctx0, Qcur, 0, 2, 1, 3); - cb(q_trans, "q_trans", il); - - qkv = ggml_mul_mat(ctx0, kv_new, q_trans); - cb(qkv, "qkv", il); - } else if(n_seq_tokens > 1) { - // lightning attention - general multi token case for PP + { + // lightning attention ggml_tensor * q_decay_exp = la->inp_q_decay; ggml_tensor * k_decay_exp = la->inp_k_decay; diff --git a/src/models/models.h b/src/models/models.h index 969429e3b6f..9b87a40d5af 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -6,6 +6,9 @@ // note: almost all graphs require at least sqrtf, so include cmath globally #include +#include + +class llama_memory_hybrid_idx_context; // // base classes @@ -2272,6 +2275,111 @@ struct llama_model_qwen35 : public llama_model_base { }; +struct llama_model_qwen4exp : public llama_model_base { + llama_model_qwen4exp(const struct llama_model_params & params) : llama_model_base(params) {} + + class llm_graph_input_qsa; + + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + struct graph : public llm_build_delta_net_base { + graph(const llama_model & model, const llm_graph_params & params); + private: + // HC replaces every layer norm: residual is [n_embd, hc, n_tokens] + ggml_tensor * build_hc_mix( + ggml_tensor * x, + ggml_tensor * w_norm, + ggml_tensor * w_down, + ggml_tensor * w_up, + ggml_tensor * w_inject, + ggml_tensor ** inject, + int il); + + ggml_tensor * build_hc_combine( + ggml_tensor * residual, + ggml_tensor * block_out, + ggml_tensor * inject, + int il); + + ggml_tensor * build_layer_attn( + llm_graph_input_attn_kv * inp_attn, + const llama_memory_hybrid_idx_context * mctx_hyb, + ggml_tensor * cur, + ggml_tensor * inp_pos, + int * sections, + int il); + + // dense self-attention restricted to the cells that top_k names + ggml_tensor * build_attn_qsa( + llm_graph_input_attn_kv * inp, + ggml_tensor * q_cur, + ggml_tensor * k_cur, + ggml_tensor * v_cur, + ggml_tensor * top_k, + float kq_scale, + int il); + + // the QSA cache layout inputs do not depend on the layer, only on its compress ratio, + // so the layers sharing a ratio share one input set + std::map qsa_inps; + + // QSA: token indices this layer's queries may attend to, or nullptr for dense + ggml_tensor * build_qsa_top_k( + const llama_memory_hybrid_idx_context * mctx_hyb, + ggml_tensor * cur, + ggml_tensor * inp_pos, + ggml_tensor * kq_mask, + int * sections, + int il); + + ggml_tensor * build_layer_attn_linear( + llm_graph_input_rs * inp, + ggml_tensor * cur, + int il); + + ggml_tensor * build_layer_ffn( + ggml_tensor * cur, + int il); + + ggml_tensor * build_norm_gated( + ggml_tensor * input, + ggml_tensor * weights, + ggml_tensor * gate, + int layer); + + // build_rs writes the state tensor in place, so one gather per cache tensor is reused + std::map rs_rows; + + // one conv history per cache tensor: delta-net and PLE each have their own + ggml_tensor * build_conv_state_at( + llm_graph_input_rs * inp, + ggml_tensor * conv_states_all, + ggml_tensor * x, + int64_t state_cols, + int64_t channels, + int il); + + ggml_tensor * build_inp_ple( + const llama_memory_hybrid_idx_context * mctx_hyb); + + ggml_tensor * build_ple( + llm_graph_input_rs * inp, + ggml_tensor * emb, + ggml_tensor * hidden, + int il); + + // returns pair of qkv, z + std::pair build_qkvz( + ggml_tensor * input, + int il); + + const llama_model & model; + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + struct llama_model_qwen35moe : public llama_model_base { llama_model_qwen35moe(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/src/models/nanbeige.cpp b/src/models/nanbeige.cpp index 3a546600fa2..7d5a6bbdb6b 100644 --- a/src/models/nanbeige.cpp +++ b/src/models/nanbeige.cpp @@ -103,6 +103,7 @@ llama_model_nanbeige::graph::graph(const llama_model & model, const llm_graph_pa ggml_tensor * inp_out_ids = build_inp_out_ids(); for (int il = 0; il < n_layer; ++il) { + res->t_layer_inp[il] = inpL; ggml_tensor * inpSA = inpL; cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il); diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp new file mode 100644 index 00000000000..abf6a0502fb --- /dev/null +++ b/src/models/qwen4exp.cpp @@ -0,0 +1,1213 @@ +#include "models.h" +#include "llama-impl.h" +#include "llama-memory-hybrid-idx.h" +#include "llama-memory-recurrent.h" + +#include +#include + +void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false); + ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false); + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + + ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, true); + + ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv); + ml.get_key(LLM_KV_SSM_INNER_SIZE, hparams.ssm_d_inner); + ml.get_key(LLM_KV_SSM_STATE_SIZE, hparams.ssm_d_state); + ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank); + ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group); + GGML_ASSERT(hparams.ssm_d_conv > 0 && hparams.ssm_d_inner > 0 && hparams.ssm_d_state > 0 && + hparams.ssm_dt_rank > 0 && hparams.ssm_n_group > 0); + + // HC; low_rank is qwen4exp-specific, DeepSeek-V4 leaves it absent (full rank) + ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult); + ml.get_key(LLM_KV_HYPER_CONNECTION_LOW_RANK, hparams.hc_low_rank); + GGML_ASSERT(hparams.dsv4_hc_mult > 0 && hparams.hc_low_rank > 0); + hparams.n_embd_out_impl = hparams.dsv4_hc_mult * hparams.n_embd; + + ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head); + ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size); + ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); + GGML_ASSERT(hparams.indexer_n_head > 0 + && hparams.indexer_head_size > 0 + && hparams.indexer_top_k > 0); + ml.get_key_or_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, hparams.n_layer_all, false); + + // PLE n-gram hash embeddings; if the key group is absent every field stays zero + hparams.is_ple_impl.reset(); + hparams.ple_n_heads = 0; + + uint32_t n_ple = 0; + ml.get_arr_n(LLM_KV_PLE_LAYERS, n_ple, false); + if (n_ple > 0) { + std::vector ple_layers; + ml.get_arr(LLM_KV_PLE_LAYERS, ple_layers); + GGML_ASSERT(n_ple == 1 && "qwen4exp supports only one PLE layer"); + for (uint32_t il : ple_layers) { + if (il >= hparams.n_layer_all) { + throw std::runtime_error(format("PLE layer %u is out of range", il)); + } + hparams.is_ple_impl.set(il); + } + + ml.get_key(LLM_KV_PLE_NGRAM_SIZE, hparams.ple_ngram_size); + ml.get_key(LLM_KV_PLE_HEADS_PER_NGRAM, hparams.ple_heads_per_ngram); + ml.get_key(LLM_KV_PLE_CONV_KERNEL, hparams.ple_conv_kernel); + ml.get_key(LLM_KV_PLE_EOS_TOKEN_ID, hparams.ple_eos_token_id); + // optional: files written before this key fall back to the EOS token + ml.get_key(LLM_KV_PLE_IMAGE_TOKEN_ID, hparams.ple_image_token_id, false); + ml.get_key(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.n_embd_per_layer); + GGML_ASSERT(hparams.ple_conv_kernel > 0 && hparams.n_embd_per_layer > 0); + + hparams.ple_n_heads = (hparams.ple_ngram_size - 1) * hparams.ple_heads_per_ngram; + hparams.ple_head_dim = hparams.n_embd_per_layer; + if (hparams.ple_ngram_size < 2 || hparams.ple_ngram_size > LLAMA_MAX_PLE_NGRAM) { + throw std::runtime_error(format("PLE n-gram size %u is out of range", hparams.ple_ngram_size)); + } + if (hparams.ple_n_heads == 0 || hparams.ple_n_heads > LLAMA_MAX_PLE_HEADS) { + throw std::runtime_error(format("PLE head count %u is out of range", hparams.ple_n_heads)); + } + + ml.get_arr(LLM_KV_PLE_LAYER_MULTIPLIERS, hparams.ple_layer_multipliers); + + // the file stores the head ranges as uint64, so read at that width and narrow to the int32 the gather uses + std::array head_offsets = {}; + std::array head_vocab_sizes = {}; + ml.get_arr(LLM_KV_PLE_HEAD_OFFSETS, head_offsets); + ml.get_arr(LLM_KV_PLE_HEAD_VOCAB_SIZES, head_vocab_sizes); + for (uint32_t h = 0; h < hparams.ple_n_heads; ++h) { + if (head_vocab_sizes[h] == 0 || + head_offsets[h] > INT32_MAX || + head_vocab_sizes[h] > INT32_MAX || + head_offsets[h] + head_vocab_sizes[h] > INT32_MAX) { + throw std::runtime_error(format("PLE head %u range does not fit the int32 row index", h)); + } + hparams.ple_head_offsets[h] = (uint32_t) head_offsets[h]; + hparams.ple_head_vocab_sizes[h] = (uint32_t) head_vocab_sizes[h]; + } + } + + // linear attention everywhere except every full_attention_interval-th layer + if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) { + uint32_t full_attn_interval = 4; + ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false); + GGML_ASSERT(full_attn_interval > 0); + for (uint32_t i = 0; i < hparams.n_layer_all; ++i) { + hparams.is_recr_impl[i] = (i < hparams.n_layer()) && ((i + 1) % full_attn_interval != 0); + } + } + + switch (hparams.n_layer()) { + case 48: type = LLM_TYPE_A3B; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) { + LLAMA_LOAD_LOCALS; + + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t hc_dim = hc * n_embd; + const int64_t hc_lr = hparams.hc_low_rank; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0); + + // there is no output_norm: the final hyper-connection mixer carries it + hc_head_norm = create_tensor(tn(LLM_TENSOR_HC_HEAD_NORM, "weight"), { hc_dim }, 0); + hc_head_down = create_tensor(tn(LLM_TENSOR_HC_HEAD_DOWN, "weight"), { hc_dim, hc_lr }, 0); + hc_head_up = create_tensor(tn(LLM_TENSOR_HC_HEAD_UP, "weight"), { hc_lr, hc_dim }, 0); + + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED); + if (output == NULL) { + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED); + } + + // flat [ple_head_dim, n_rows] gather target; n_rows is padded, so read it back + if (hparams.ple_n_heads > 0) { + const std::string ple_name = tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight").str(); + const auto & ple_w = ml.require_weight(ple_name.c_str()); + const int64_t ple_rows = ple_w.tensor->ne[1]; + + // sanity check + for (uint32_t h = 0; h < hparams.ple_n_heads; ++h) { + if ((int64_t) hparams.ple_head_offsets[h] + hparams.ple_head_vocab_sizes[h] > ple_rows) { + throw std::runtime_error(format("PLE head %u range exceeds the %" PRId64 " table rows", h, ple_rows)); + } + } + per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), + { hparams.ple_head_dim, ple_rows }, TENSOR_READ_LAZY); + } + + for (int il = 0; il < n_layer; ++il) { + auto & layer = layers[il]; + + const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used; + const int64_t n_ff_shexp = hparams.n_ff_shexp ? hparams.n_ff_shexp : n_ff; + + const int64_t head_k_dim = hparams.ssm_d_state; + const int64_t head_v_dim = hparams.ssm_d_state; + const int64_t n_k_heads = hparams.ssm_n_group; + const int64_t n_v_heads = hparams.ssm_dt_rank; + const int64_t key_dim = head_k_dim * n_k_heads; + const int64_t value_dim = head_v_dim * n_v_heads; + const int64_t conv_dim = key_dim * 2 + value_dim; + + // two HC modules per layer: before the token mixer, before the MoE + layer.hc_attn_norm = create_tensor(tn(LLM_TENSOR_HC_ATTN_NORM, "weight", il), { hc_dim }, 0); + layer.hc_attn_down = create_tensor(tn(LLM_TENSOR_HC_ATTN_DOWN, "weight", il), { hc_dim, hc_lr }, 0); + layer.hc_attn_up = create_tensor(tn(LLM_TENSOR_HC_ATTN_UP, "weight", il), { hc_lr, hc_dim }, 0); + layer.hc_attn_inject = create_tensor(tn(LLM_TENSOR_HC_ATTN_INJECT, "weight", il), { hc_dim, hc }, 0); + layer.hc_ffn_norm = create_tensor(tn(LLM_TENSOR_HC_FFN_NORM, "weight", il), { hc_dim }, 0); + layer.hc_ffn_down = create_tensor(tn(LLM_TENSOR_HC_FFN_DOWN, "weight", il), { hc_dim, hc_lr }, 0); + layer.hc_ffn_up = create_tensor(tn(LLM_TENSOR_HC_FFN_UP, "weight", il), { hc_lr, hc_dim }, 0); + layer.hc_ffn_inject = create_tensor(tn(LLM_TENSOR_HC_FFN_INJECT, "weight", il), { hc_dim, hc }, 0); + + if (!hparams.is_recr(il)) { + // full attention: wq holds [q|gate] interleaved per head + create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, 0); + + layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, 0); + layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, 0); + + const int64_t idx_dim = hparams.indexer_head_size; + layer.index_q_proj = create_tensor(tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", il), { n_embd, hparams.indexer_n_head * idx_dim }, 0); + layer.index_k_proj = create_tensor(tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", il), { n_embd, idx_dim }, 0); + layer.index_q_norm = create_tensor(tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", il), { idx_dim }, 0); + layer.index_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", il), { idx_dim }, 0); + } else { + layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", il), { n_embd, key_dim * 2 + value_dim }, 0); + layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, value_dim }, 0); + layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", il), { hparams.ssm_d_conv, conv_dim }, 0); + layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { hparams.ssm_dt_rank }, 0); + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, il), { hparams.ssm_dt_rank }, 0); + layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", il), { n_embd, n_v_heads }, 0); + layer.ssm_alpha = create_tensor(tn(LLM_TENSOR_SSM_ALPHA, "weight", il), { n_embd, n_v_heads }, 0); + layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), { head_v_dim }, 0); + layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", il), { value_dim, n_embd }, 0); + } + + if (hparams.is_ple(il)) { + layer.ple_key = create_tensor(tn(LLM_TENSOR_PLE_KEY, "weight", il), { n_embd, hc_dim }, 0); + layer.ple_value = create_tensor(tn(LLM_TENSOR_PLE_VALUE, "weight", il), { n_embd, n_embd }, 0); + layer.ple_norm_key = create_tensor(tn(LLM_TENSOR_PLE_NORM_KEY, "weight", il), { hc_dim }, 0); + layer.ple_norm_query = create_tensor(tn(LLM_TENSOR_PLE_NORM_QUERY, "weight", il), { hc_dim }, 0); + layer.ple_norm_conv = create_tensor(tn(LLM_TENSOR_PLE_NORM_CONV, "weight", il), { hc_dim }, 0); + layer.ple_conv1d = create_tensor(tn(LLM_TENSOR_PLE_CONV1D, "weight", il), { hparams.ple_conv_kernel, hc_dim }, 0); + } + + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, 0); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, 0); + create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, 0); + + layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, 0); + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, 0); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_shexp }, 0); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_shexp, n_embd }, 0); + } +} + +std::unique_ptr llama_model_qwen4exp::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +// Hyper-connections keep hc parallel residual streams [n_embd, hc, T] in place of layer norms. +// Returns the mixed [n_embd, T] stream; `inject` gets the [hc, T] scatter weights. +ggml_tensor * llama_model_qwen4exp::graph::build_hc_mix( + ggml_tensor * x, + ggml_tensor * w_norm, + ggml_tensor * w_down, + ggml_tensor * w_up, + ggml_tensor * w_inject, + ggml_tensor ** inject, + int il) { + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t hc_dim = hc * n_embd; + const int64_t nt = x->ne[2]; + + // grouped RMSNorm: reduce over one stream, then scale all streams with the [hc_dim] gamma + // the converter folded each gamma to (1 + w) + ggml_tensor * xn = ggml_rms_norm(ctx0, x, hparams.f_norm_rms_eps); + xn = ggml_reshape_2d(ctx0, xn, hc_dim, nt); + xn = ggml_mul(ctx0, xn, w_norm); + cb(xn, "hc_norm", il); + + ggml_tensor * lo = build_lora_mm(w_down, xn); + lo = ggml_silu(ctx0, ggml_scale(ctx0, lo, 1.0f / (float) hc)); + ggml_tensor * gate = ggml_sigmoid(ctx0, build_lora_mm(w_up, lo)); + cb(gate, "hc_gate", il); + + ggml_tensor * gated = ggml_mul(ctx0, xn, gate); + gated = ggml_reshape_3d(ctx0, gated, n_embd, hc, nt); + + // collapse the streams by their mean + ggml_tensor * mixed = ggml_view_2d(ctx0, gated, n_embd, nt, + ggml_row_size(gated->type, n_embd) * hc, 0); + mixed = ggml_cont(ctx0, mixed); + for (int64_t c = 1; c < hc; ++c) { + ggml_tensor * s = ggml_view_2d(ctx0, gated, n_embd, nt, + ggml_row_size(gated->type, n_embd) * hc, + ggml_row_size(gated->type, n_embd) * c); + mixed = ggml_add(ctx0, mixed, s); + } + mixed = ggml_scale(ctx0, mixed, 1.0f / (float) hc); + cb(mixed, "hc_mixed", il); + + if (inject) { + *inject = build_lora_mm(w_inject, xn); + cb(*inject, "hc_inject", il); + } + + return mixed; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_hc_combine( + ggml_tensor * residual, + ggml_tensor * block_out, + ggml_tensor * inject, + int il) { + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t nt = residual->ne[2]; + + // 2*sigmoid centres the scatter weights on 1, so a zero injection is a plain residual add + ggml_tensor * w = ggml_sigmoid(ctx0, ggml_scale(ctx0, inject, 1.0f / (float) hc)); + w = ggml_scale(ctx0, w, 2.0f); + w = ggml_reshape_3d(ctx0, w, 1, hc, nt); + + ggml_tensor * b = ggml_reshape_3d(ctx0, block_out, n_embd, 1, nt); + b = ggml_repeat_4d(ctx0, b, n_embd, hc, nt, 1); + + ggml_tensor * cur = ggml_add(ctx0, residual, ggml_mul(ctx0, b, w)); + cb(cur, "hc_combine", il); + + return cur; +} + +llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_params & params) : + llm_build_delta_net_base(params), model(model) { + const int64_t hc = hparams.dsv4_hc_mult; + + GGML_ASSERT(hparams.n_embd_head_v() == hparams.n_embd_head_k()); + + int sections[4]; + std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections); + + ggml_tensor * inpL = build_inp_embd(model.tok_embd); + cb(inpL, "model.input_embed", -1); + ggml_build_forward_expand(gf, inpL); + + auto * inp = build_inp_mem_hybrid(); + + // qwen4exp always builds llama_memory_hybrid_idx, so this downcast is safe + // the indexer cache inside it is absent when the GGUF has no indexer tensors + const auto * mctx_hyb = static_cast(inp->mctx); + + const llama_kv_cache_context * mctx_idx = mctx_hyb->get_idx(); + if (mctx_idx) { + GGML_ASSERT(mctx_idx->get_n_kv() == inp->mctx->get_attn()->get_n_kv() && + "the indexer cache must track the attention cache cell for cell"); + } + + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + ggml_tensor * ple_emb = nullptr; + if (hparams.ple_n_heads > 0) { + ple_emb = build_inp_ple(mctx_hyb); + // make sure ple_emb and build_inp_embd are in the same graph split + ggml_build_forward_expand(gf, ple_emb); + } + + // the wide residual starts as hc identical copies of the embedding + ggml_tensor * res_hc = ggml_repeat_4d(ctx0, + ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens), + n_embd, hc, n_tokens, 1); + cb(res_hc, "hc_init", -1); + + for (int il = 0; il < n_layer; ++il) { + res->t_layer_inp[il] = res_hc; + + if (hparams.is_ple(il)) { + res_hc = build_ple(inp->get_recr(), ple_emb, res_hc, il); + } + + ggml_tensor * inject = nullptr; + ggml_tensor * cur = build_hc_mix(res_hc, + model.layers[il].hc_attn_norm, + model.layers[il].hc_attn_down, + model.layers[il].hc_attn_up, + model.layers[il].hc_attn_inject, + &inject, il); + + ggml_build_forward_expand(gf, cur); + + if (hparams.is_recr(il)) { + cur = build_layer_attn_linear(inp->get_recr(), cur, il); + } else { + cur = build_layer_attn(inp->get_attn(), mctx_hyb, cur, inp_pos, sections, il); + } + + if (il == n_layer - 1 && inp_out_ids) { + // everything below is per token, so drop the rows that produce no output + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inject = ggml_get_rows(ctx0, inject, inp_out_ids); + + res_hc = ggml_reshape_2d(ctx0, res_hc, n_embd*hc, res_hc->ne[2]); + res_hc = ggml_get_rows(ctx0, res_hc, inp_out_ids); + res_hc = ggml_reshape_3d(ctx0, res_hc, n_embd, hc, res_hc->ne[1]); + } + + res_hc = build_hc_combine(res_hc, cur, inject, il); + + cur = build_hc_mix(res_hc, + model.layers[il].hc_ffn_norm, + model.layers[il].hc_ffn_down, + model.layers[il].hc_ffn_up, + model.layers[il].hc_ffn_inject, + &inject, il); + + cur = build_layer_ffn(cur, il); + cb(cur, "ffn_out", il); + + res_hc = build_hc_combine(res_hc, cur, inject, il); + + // "l_last" is the layer output name that build_cvec and imatrix look for + cb(res_hc, "l_last", il); + } + + // the final mixer is the output norm: there is no separate one + ggml_tensor * cur = build_hc_mix(res_hc, + model.hc_head_norm, model.hc_head_down, model.hc_head_up, + nullptr, nullptr, -1); + + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = build_lora_mm(model.output, cur, model.output_s); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} + +std::pair llama_model_qwen4exp::graph::build_qkvz( + ggml_tensor * input, + int il) { + const int64_t n_seqs = ubatch.n_seqs; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + + ggml_tensor * qkv_mixed = build_lora_mm(model.layers[il].wqkv, input, model.layers[il].wqkv_s); + qkv_mixed = ggml_reshape_3d(ctx0, qkv_mixed, qkv_mixed->ne[0], n_seq_tokens, n_seqs); + cb(qkv_mixed, "linear_attn_qkv_mixed", il); + + ggml_tensor * z = build_lora_mm(model.layers[il].wqkv_gate, input, model.layers[il].wqkv_gate_s); + cb(z, "z", il); + + return { qkv_mixed, z }; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_norm_gated( + ggml_tensor * input, + ggml_tensor * weights, + ggml_tensor * gate, + int layer) { + // the one numerical difference from Qwen3.5's GDN: sigmoid output gate, not silu + ggml_tensor * normalized = build_norm(input, weights, nullptr, LLM_NORM_RMS, layer); + ggml_tensor * gated = ggml_sigmoid(ctx0, gate); + + return ggml_mul(ctx0, normalized, gated); +} + +// QSA attends to a budget of whole blocks of compress_ratio tokens, plus the incomplete tail +// one mean-pooled indexer key scores each block; set_input resolves the cache layout +class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i { +public: + llm_graph_input_qsa(const llama_memory_hybrid_idx_context * mctx, uint32_t ratio, bool blk_bias) : + mctx(mctx), ratio(ratio), blk_bias(blk_bias) {} + virtual ~llm_graph_input_qsa() = default; + + void set_input(const llama_ubatch * ubatch) override { + mctx->get_idx()->set_input_k_idxs(k_idxs, ubatch); + mctx->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio, blk_bias); + } + + bool can_reuse(const llm_graph_params & params) override { + mctx = static_cast(params.mctx); + + const auto * idx = mctx->get_idx(); + if (idx == nullptr) { + return false; + } + + const int64_t n_kv = idx->get_n_kv(); + const int64_t n_stream = mctx->get_n_stream(); + const int64_t n_blocks = (n_kv + ratio - 1)/ratio; + + bool res = true; + + res &= params.ubatch.n_tokens % n_stream == 0; + + res &= k_idxs->ne[0] == params.ubatch.n_tokens; + res &= cell_blk->ne[0] == n_kv; + res &= cell_blk->ne[1] == n_stream; + res &= blk_cells->ne[0] == (int64_t) ratio*n_blocks; + res &= blk_pos->ne[0] == 4*n_blocks*n_stream; + res &= bias->ne[0] == (blk_bias ? n_blocks : n_kv); + res &= bias->ne[1] == params.ubatch.n_tokens/n_stream; + + return res; + } + + // per stream: a cell index names a different token in each stream + ggml_tensor * k_idxs = nullptr; // I32 [n_tokens] + ggml_tensor * cell_blk = nullptr; // I32 [n_kv, n_stream] + ggml_tensor * blk_cells = nullptr; // I32 [ratio*n_blocks, n_stream] + ggml_tensor * blk_pos = nullptr; // I32 [4*n_blocks*n_stream] + ggml_tensor * bias = nullptr; // F32 [n_blocks or n_kv, n_tokens/n_stream, n_stream] + + const llama_memory_hybrid_idx_context * mctx; + const uint32_t ratio; + + // the per-cell half of the bias is the attention mask, so only the per-block half is uploaded + const bool blk_bias; +}; + +ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( + const llama_memory_hybrid_idx_context * mctx_hyb, + ggml_tensor * cur, + ggml_tensor * inp_pos, + ggml_tensor * kq_mask, + int * sections, + int il) { + const llama_kv_cache_context * mctx_idx = mctx_hyb->get_idx(); + + const int64_t idx_dim = hparams.indexer_head_size; + const int64_t n_idx_h = hparams.indexer_n_head; + const int64_t r = hparams.dsv4_compress_ratios[il]; + const int64_t n_kv = mctx_idx->get_n_kv(); + + GGML_ASSERT(r > 0); + + const int64_t n_blocks = (n_kv + r - 1)/r; + + // build_attn_qsa and the KQ mask need the tokens to divide evenly across the streams + const int64_t n_stream = mctx_hyb->get_n_stream(); + GGML_ASSERT(n_tokens % n_stream == 0); + const int64_t n_tps = n_tokens/n_stream; + + // only the "which block is visible" half of the bias varies per block + // the rest is the visible/not test the attention mask already carries, so upload the per-block half only: 1/ratio of the cells + // alibi writes distances instead of a mask and non-causal keeps future cells, so both opt out + // the mask also holds an mrope rule for the query's own position, but only 2d image positions can differ there + const bool blk_bias = kq_mask != nullptr && + kq_mask->ne[0] == n_kv && kq_mask->ne[1] == n_tps && kq_mask->ne[3] == n_stream && + cparams.causal_attn && !hparams.use_alibi; + + // nothing above depends on the layer, so the layers sharing a ratio share one input set + llm_graph_input_qsa * inp = nullptr; + + const auto it = qsa_inps.find((uint32_t) r); + if (it != qsa_inps.end()) { + inp = it->second; + } else { + auto qsa = std::make_unique(mctx_hyb, (uint32_t) r, blk_bias); + + qsa->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch); + qsa->cell_blk = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, n_stream); + qsa->blk_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, r*n_blocks, n_stream); + qsa->blk_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 4*n_blocks*n_stream); + qsa->bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, blk_bias ? n_blocks : n_kv, n_tps, n_stream); + + ggml_set_input(qsa->cell_blk); + ggml_set_input(qsa->blk_cells); + ggml_set_input(qsa->blk_pos); + ggml_set_input(qsa->bias); + + inp = qsa.get(); + res->add_input(std::move(qsa)); + qsa_inps.emplace((uint32_t) r, inp); + } + + // cached indexer keys are raw: pooling precedes norm and rotation, so apply neither + ggml_tensor * k_raw = build_lora_mm(model.layers[il].index_k_proj, cur); + k_raw = ggml_reshape_3d(ctx0, k_raw, idx_dim, 1, n_tokens); + cb(k_raw, "indexer_k_raw", il); + + ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, k_raw, inp->k_idxs, il)); + + // one key head, so rows are contiguous. get_k gives [idx_dim, n_head_kv, n_kv, n_stream]. + ggml_tensor * k_all = mctx_idx->get_k(ctx0, il); + k_all = ggml_view_3d(ctx0, k_all, idx_dim, n_kv, n_stream, k_all->nb[2], k_all->nb[3], 0); + + // gathers per stream: blk_cells row s indexes stream s's own cells + ggml_tensor * members = ggml_get_rows(ctx0, k_all, inp->blk_cells); + members = ggml_reshape_4d(ctx0, members, idx_dim, r, n_blocks, n_stream); + + // mean over the block members; r is small, so summing slices beats a transpose plus sum_rows + ggml_tensor * pooled = nullptr; + for (int64_t i = 0; i < r; ++i) { + ggml_tensor * slice = ggml_cont(ctx0, + ggml_view_3d(ctx0, members, idx_dim, n_blocks, n_stream, + members->nb[2], members->nb[3], i*members->nb[1])); + pooled = pooled ? ggml_add(ctx0, pooled, slice) : slice; + } + pooled = ggml_scale(ctx0, pooled, 1.0f/(float) r); + cb(pooled, "indexer_k_pooled", il); + + // rope wants [n_dims, n_head, n_tokens]: lay every stream's blocks flat, split after. + pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, 1, n_blocks*n_stream); + pooled = build_norm(pooled, model.layers[il].index_k_norm, nullptr, LLM_NORM_RMS, il); + pooled = ggml_rope_multi(ctx0, pooled, inp->blk_pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, n_blocks, n_stream); + cb(pooled, "indexer_k", il); + + ggml_tensor * q = build_lora_mm(model.layers[il].index_q_proj, cur); + q = ggml_reshape_3d(ctx0, q, idx_dim, n_idx_h, n_tokens); + q = build_norm(q, model.layers[il].index_q_norm, nullptr, LLM_NORM_RMS, il); + q = ggml_rope_multi(ctx0, q, inp_pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(q, "indexer_q", il); + + // rectify each head dot product before the sum, as in the DeepSeek lightning indexer + // mul_mat matches ne[2], so the queries of stream s only meet the blocks of stream s + ggml_tensor * score = ggml_mul_mat(ctx0, pooled, + ggml_reshape_3d(ctx0, ggml_cont(ctx0, q), idx_dim, n_idx_h*n_tps, n_stream)); + score = ggml_reshape_4d(ctx0, score, n_blocks, n_idx_h, n_tps, n_stream); + score = ggml_relu(ctx0, score); + score = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)); + score = ggml_sum_rows(ctx0, score); + score = ggml_reshape_3d(ctx0, score, n_blocks, n_tps, n_stream); + cb(score, "indexer_score", il); + + // one value per block, so it is cheaper to bias here than after the cells are expanded + if (blk_bias) { + score = ggml_add(ctx0, score, inp->bias); + } + + // every token of a block gets the block score; the budget is whole blocks, so top-k cuts on a block boundary + ggml_tensor * expanded = ggml_get_rows(ctx0, + ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)), inp->cell_blk); + expanded = ggml_cont(ctx0, ggml_permute(ctx0, expanded, 1, 0, 2, 3)); + + if (blk_bias) { + // flash attention keeps the mask in f16; the scores are f32 + ggml_tensor * mask = kq_mask->type == GGML_TYPE_F32 ? kq_mask : ggml_cast(ctx0, kq_mask, GGML_TYPE_F32); + expanded = ggml_add(ctx0, expanded, ggml_reshape_3d(ctx0, mask, n_kv, n_tps, n_stream)); + } else { + expanded = ggml_add(ctx0, expanded, inp->bias); + } + cb(expanded, "indexer_score_tokens", il); + + // the reference returns indexer_top_k + compress_ratio - 1: whole blocks plus the tail + const int64_t width = std::min(n_kv, (int64_t) hparams.indexer_top_k + r - 1); + + ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, expanded, width)); + + // build_attn_qsa reads [n_top_k, n_batch, 1, n_stream], matching the KQ mask. + top_k = ggml_reshape_4d(ctx0, top_k, width, n_tps, 1, n_stream); + cb(top_k, "indexer_top_k", il); + + return top_k; +} + +// Dense GQA self-attention restricted to the cells that top_k names. +// The mask build below copies the MLA sparse path in llm_graph_context::build_attn. +ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( + llm_graph_input_attn_kv * inp, + ggml_tensor * q_cur, + ggml_tensor * k_cur, + ggml_tensor * v_cur, + ggml_tensor * top_k, + float kq_scale, + int il) { + // rotate q/k/v before they reach a quantized cache, as the dense path does. the indexer + // has already scored with its own query in build_qsa_top_k, so top_k is unaffected. + if (inp->self_k_rot) { + q_cur = llama_mul_mat_hadamard(ctx0, q_cur, inp->self_k_rot); + k_cur = llama_mul_mat_hadamard(ctx0, k_cur, inp->self_k_rot); + } + + if (inp->self_v_rot) { + v_cur = llama_mul_mat_hadamard(ctx0, v_cur, inp->self_v_rot); + } + + // these nodes are added to the graph together so that they are not reordered + // by doing so, the number of splits in the graph is reduced + // expand k later to enable rope fusion which directly writes into k-v cache + ggml_build_forward_expand(gf, q_cur); + ggml_build_forward_expand(gf, v_cur); + ggml_build_forward_expand(gf, k_cur); + + const auto * mctx_cur = inp->mctx; + + // store to KV cache + { + const auto & k_idxs = inp->get_k_idxs(); + const auto & v_idxs = inp->get_v_idxs(); + + ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); + ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il)); + } + + ggml_tensor * kq_mask = inp->get_kq_mask(); + + // prepare new kq mask - starts filled with -INFINITY + ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY); + + // reshape KQ mask into tensor with rows of size 1: + // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream] + kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0); + + // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1] + ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0); + + // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream] + // this will be our source of zero values for unmasking top k mask elements + ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); + zeros = ggml_fill(ctx0, zeros, 0.0f); + + // modify KQ mask by unmasking elements that are in top_k indices + // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1]) + ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d); + + // reshape to restore the original shape of KQ mask: + // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream] + kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0); + + // combine with the original kq mask + kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask); + + ggml_tensor * q = q_cur; + ggml_tensor * k = mctx_cur->get_k(ctx0, il); + ggml_tensor * v = mctx_cur->get_v(ctx0, il); + + ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, kq_scale, il); + cb(cur, "kqv_out", il); + + // the rotation is its own inverse, so undo it on the value side of the output + if (inp->self_v_rot) { + cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot); + } + + return cur; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( + llm_graph_input_attn_kv * inp, + const llama_memory_hybrid_idx_context * mctx_hyb, + ggml_tensor * cur, + ggml_tensor * inp_pos, + int * sections, + int il) { + const int64_t n_embd_head = hparams.n_embd_head_v(); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + + // indexer reads the same block input as q/k/v; no cache or no ratio means dense + const bool qsa = mctx_hyb->get_idx() != nullptr && hparams.dsv4_compress_ratios[il] > 0; + + ggml_tensor * top_k = qsa ? build_qsa_top_k(mctx_hyb, cur, inp_pos, inp->get_kq_mask(), sections, il) : nullptr; + + // Qwen3Next uses a single Q projection that outputs query + gate + ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); // [ (n_embd_head * 2) * n_head, n_tokens ] + cb(Qcur_full, "Qcur_full", il); + + ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, + ggml_element_size(Qcur_full) * n_embd_head * 2, + ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head, 0); + cb(Qcur, "Qcur_reshaped", il); + + Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il); + cb(Qcur, "Qcur_normed", il); + + ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur, model.layers[il].wk_s); + cb(Kcur, "Kcur", il); + + ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur, model.layers[il].wv_s); + cb(Vcur, "Vcur", il); + + Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); + Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il); + cb(Kcur, "Kcur_normed", il); + + ggml_tensor * gate = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, + ggml_element_size(Qcur_full) * n_embd_head * 2, + ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head, + ggml_element_size(Qcur_full) * n_embd_head); + gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens); + cb(gate, "gate_reshaped", il); + + Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); + + // Apply IMRoPE + Qcur = ggml_rope_multi( + ctx0, Qcur, inp_pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + + Kcur = ggml_rope_multi( + ctx0, Kcur, inp_pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + + const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale; + + if (top_k) { + cur = build_attn_qsa(inp, Qcur, Kcur, Vcur, top_k, kq_scale, il); + } else { + cur = build_attn(inp, + nullptr, nullptr, nullptr, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + } + cb(cur, "attn_pregate", il); + + ggml_tensor * gate_sigmoid = ggml_sigmoid(ctx0, gate); + cb(gate_sigmoid, "gate_sigmoid", il); + + cur = ggml_mul(ctx0, cur, gate_sigmoid); + cb(cur, "attn_gated", il); + + cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s); + cb(cur, "attn_output", il); + + return cur; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn_linear( + llm_graph_input_rs * inp, + ggml_tensor * cur, + int il) { + const auto * mctx_cur = inp->mctx; + + const int64_t d_inner = hparams.ssm_d_inner; + const int64_t n_seqs = ubatch.n_seqs; + const int64_t head_k_dim = hparams.ssm_d_state; + const int64_t num_k_heads = hparams.ssm_n_group; + const int64_t num_v_heads = hparams.ssm_dt_rank; + const int64_t head_v_dim = hparams.ssm_d_state; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + + GGML_ASSERT(n_seqs != 0); + GGML_ASSERT(ubatch.equal_seqs()); + GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs); + GGML_ASSERT(head_v_dim * num_v_heads == d_inner); + + auto qkvz = build_qkvz(cur, il); + ggml_tensor * qkv_mixed = qkvz.first; + ggml_tensor * z = qkvz.second; + + ggml_tensor * beta = build_lora_mm(model.layers[il].ssm_beta, cur, model.layers[il].ssm_beta_s); + beta = ggml_reshape_4d(ctx0, beta, 1, num_v_heads, n_seq_tokens, n_seqs); + cb(beta, "beta", il); + + beta = ggml_sigmoid(ctx0, beta); + cb(beta, "beta_sigmoid", il); + + ggml_tensor * alpha = build_lora_mm(model.layers[il].ssm_alpha, cur, model.layers[il].ssm_alpha_s); + alpha = ggml_reshape_3d(ctx0, alpha, num_v_heads, n_seq_tokens, n_seqs); + cb(alpha, "alpha", il); + + ggml_tensor * alpha_biased = ggml_add(ctx0, alpha, model.layers[il].ssm_dt); + ggml_tensor * alpha_softplus = ggml_softplus(ctx0, alpha_biased); + cb(alpha_softplus, "a_softplus", il); + + ggml_tensor * gate = ggml_mul(ctx0, alpha_softplus, model.layers[il].ssm_a); // -A_log.exp() * softplus + cb(gate, "gate", il); + + gate = ggml_reshape_4d(ctx0, gate, 1, num_v_heads, n_seq_tokens, n_seqs); + + ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); + ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il); + + ggml_tensor * conv_kernel = model.layers[il].ssm_conv1d; + const int64_t conv_kernel_size = conv_kernel->ne[0]; + + // the channels must match how load_arch_tensors sizes wqkv, not ssm_d_inner + const int64_t conv_channels = head_k_dim * num_k_heads * 2 + head_v_dim * num_v_heads; + + ggml_tensor * conv_input = build_conv_state_at(inp, conv_states_all, qkv_mixed, + conv_kernel_size - 1, conv_channels, il); + + ggml_tensor * state = build_rs(inp, ssm_states_all, hparams.n_embd_s(), n_seqs); + state = ggml_reshape_4d(ctx0, state, head_v_dim, head_v_dim, num_v_heads, n_seqs); + cb(state, "state_predelta", il); + + ggml_tensor * conv_output_proper = ggml_ssm_conv(ctx0, conv_input, conv_kernel); + cb(conv_output_proper, "conv_output_raw", il); + + ggml_tensor * conv_output_silu = ggml_silu(ctx0, conv_output_proper); + cb(conv_output_silu, "conv_output_silu", il); + + ggml_tensor * conv_qkv_mix = conv_output_silu; + + int64_t nb1_qkv = ggml_row_size(conv_qkv_mix->type, conv_channels); + + // Extract the convolved Q, K, V from conv_output + ggml_tensor * q_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_k_dim, num_k_heads, n_seq_tokens, n_seqs, + ggml_row_size(conv_qkv_mix->type, head_k_dim), + nb1_qkv, + nb1_qkv * n_seq_tokens, + 0); + + ggml_tensor * k_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_k_dim, num_k_heads, n_seq_tokens, n_seqs, + ggml_row_size(conv_qkv_mix->type, head_k_dim), + nb1_qkv, + nb1_qkv * n_seq_tokens, + head_k_dim * num_k_heads * ggml_element_size(conv_qkv_mix)); + + ggml_tensor * v_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_v_dim, num_v_heads, n_seq_tokens, n_seqs, + ggml_row_size(conv_qkv_mix->type, head_v_dim), + nb1_qkv, + nb1_qkv * n_seq_tokens, + ggml_row_size(conv_qkv_mix->type, 2 * head_k_dim * num_k_heads)); + + cb(q_conv, "q_conv", il); + cb(k_conv, "k_conv", il); + cb(v_conv, "v_conv", il); + + const float eps_norm = hparams.f_norm_rms_eps; + + q_conv = ggml_l2_norm(ctx0, q_conv, eps_norm); + k_conv = ggml_l2_norm(ctx0, k_conv, eps_norm); + + // repeat to match shapes when head keys != value keys; unneeded with the fused GDN + if (num_k_heads != num_v_heads && (!cparams.fused_gdn_ar || !cparams.fused_gdn_ch)) { + GGML_ASSERT(num_v_heads % num_k_heads == 0); + q_conv = ggml_repeat_4d(ctx0, q_conv, head_k_dim, num_v_heads, n_seq_tokens, n_seqs); + k_conv = ggml_repeat_4d(ctx0, k_conv, head_k_dim, num_v_heads, n_seq_tokens, n_seqs); + } + + cb(q_conv, "q_conv_predelta", il); + cb(k_conv, "k_conv_predelta", il); + cb(v_conv, "v_conv_predelta", il); + + ggml_tensor * output = build_recurrent_attn(inp, ssm_states_all, q_conv, k_conv, v_conv, gate, beta, state, il); + + ggml_tensor * z_2d = ggml_reshape_4d(ctx0, z, head_v_dim, num_v_heads, n_seq_tokens, n_seqs); + + // gated normalization, as self.norm(core_attn_out, z) in the reference + ggml_tensor * attn_out_norm = build_norm_gated(output, model.layers[il].ssm_norm, z_2d, il); + + ggml_tensor * final_output = ggml_reshape_3d(ctx0, attn_out_norm, head_v_dim * num_v_heads, n_seq_tokens, n_seqs); + cb(final_output, "final_output", il); + + cur = build_lora_mm(model.layers[il].ssm_out, final_output, model.layers[il].ssm_out_s); + cb(cur, "linear_attn_out", il); + + cur = ggml_reshape_2d(ctx0, cur, n_embd, n_seq_tokens * n_seqs); + + return cur; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_layer_ffn(ggml_tensor * cur, const int il) { + GGML_ASSERT(model.layers[il].ffn_gate_inp != nullptr); + + ggml_tensor * moe_out = + build_moe_ffn(cur, + model.layers[il].ffn_gate_inp, + model.layers[il].ffn_up_exps, + model.layers[il].ffn_gate_exps, + model.layers[il].ffn_down_exps, + nullptr, + n_expert, n_expert_used, + LLM_FFN_SILU, true, + hparams.expert_weights_scale, + LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, il, + nullptr, model.layers[il].ffn_gate_up_exps, + model.layers[il].ffn_up_exps_s, + model.layers[il].ffn_gate_exps_s, + model.layers[il].ffn_down_exps_s); + cb(moe_out, "ffn_moe_out", il); + + // shared experts, as in the Qwen3Next reference + if (model.layers[il].ffn_up_shexp != nullptr) { + ggml_tensor * ffn_shexp = + build_ffn(cur, + model.layers[il].ffn_up_shexp, NULL, model.layers[il].ffn_up_shexp_s, + model.layers[il].ffn_gate_shexp, NULL, model.layers[il].ffn_gate_shexp_s, + model.layers[il].ffn_down_shexp, NULL, model.layers[il].ffn_down_shexp_s, + NULL, + LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(ffn_shexp, "ffn_shexp", il); + + // shared expert has its own sigmoided gate (ffn_gate_inp_shexp, one value per token) + ggml_tensor * shared_gate = build_lora_mm(model.layers[il].ffn_gate_inp_shexp, cur); + cb(shared_gate, "shared_expert_gate", il); + + shared_gate = ggml_sigmoid(ctx0, shared_gate); + cb(shared_gate, "shared_expert_gate_sigmoid", il); + + ffn_shexp = ggml_mul(ctx0, ffn_shexp, shared_gate); + cb(ffn_shexp, "ffn_shexp_gated", il); + + cur = ggml_add(ctx0, moe_out, ffn_shexp); + cb(cur, "ffn_out", il); + } else { + cur = moe_out; + } + + return cur; +} + +// PLE n-gram hash embedding: each token gathers ple_n_heads rows of a shared table. +// mixed_n = (t[p]*m[0]) ^ ... ^ (t[p-n+1]*m[n-1]); row = mixed_n % vocab[h] + offset[h] +// The hash runs host-side because ggml has no int64 and no xor. EOS resets the window. + +class llm_graph_input_ple : public llm_graph_input_i { +public: + llm_graph_input_ple(const llama_model_qwen4exp & pmodel, + const llama_kv_cache_context * mctx) : pmodel(pmodel), mctx(mctx) {} + virtual ~llm_graph_input_ple() = default; + + void set_input(const llama_ubatch * ubatch) override; + + bool can_reuse(const llm_graph_params & params) override { + mctx = static_cast(params.mctx)->get_attn(); + return rows->ne[0] == (int64_t) pmodel.hparams.ple_n_heads * params.ubatch.n_tokens; + } + + ggml_tensor * rows = nullptr; // I32 [ple_n_heads * n_tokens] + + const llama_model_qwen4exp & pmodel; + + // the predecessor tokens live in the attention KV cells (ext.tok) + const llama_kv_cache_context * mctx; + + // scratch, reused across set_input() calls + std::vector prev; +}; + +void llm_graph_input_ple::set_input(const llama_ubatch * ubatch) { + const auto & hp = pmodel.hparams; + + // an image arrives as an embd batch, so ubatch->token is null, but every position still needs a row for ggml_get_rows + // stand in the image token id that the reference hashes, or EOS if the file has no such key + // gemma3n and gemma4 do the same with a hardcoded row 0 of per_layer_token_embd. + const llama_token img_tok = hp.ple_image_token_id != 0 + ? (llama_token) hp.ple_image_token_id + : (llama_token) hp.ple_eos_token_id; + auto tok_of = [&](int64_t k) -> llama_token { + return ubatch->token ? ubatch->token[k] : img_tok; + }; + + const int64_t n_tokens = ubatch->n_tokens; + const int64_t n_gram = hp.ple_ngram_size; + const int64_t n_heads = hp.ple_n_heads; + const int64_t per_gram = hp.ple_heads_per_ngram; + const int64_t eos = hp.ple_eos_token_id; + const int64_t n_prev = n_gram - 1; + + std::vector idx(n_heads * n_tokens); + + GGML_ASSERT(mctx != nullptr); + + for (int64_t i = 0; i < n_tokens; ++i) { + // the preceding tokens would be ambiguous, see get_prev_tokens() + GGML_ASSERT(ubatch->n_seq_id[i] == 1 && "PLE n-gram embeddings do not support tokens shared by multiple sequences"); + } + + // predecessors come from the KV cells (ext.tok); apply_ubatch() already stored this ubatch, so its own tokens count too + mctx->get_prev_tokens(*ubatch, n_prev, prev); + + for (int64_t i = 0; i < n_tokens; ++i) { + // an EOS in the window resets everything at or before it + // a missing predecessor (before the sequence start, or no cached cell) reads as EOS + // the EOS of the token itself does not cut its own context, as in the reference + std::vector ctx(n_gram); + ctx[0] = tok_of(i); + bool cut = false; + for (int64_t s = 1; s < n_gram; ++s) { + // predecessor s positions back; prev[] is oldest-first, missing entries are LLAMA_TOKEN_NULL + const llama_token t = cut ? LLAMA_TOKEN_NULL : prev[i*n_prev + (n_prev - s)]; + cut = cut || t < 0 || t == eos; + ctx[s] = cut ? eos : t; + } + + for (int64_t n = 2; n <= n_gram; ++n) { + uint64_t mixed = (uint64_t) ctx[0] * hp.ple_layer_multipliers[0]; + for (int64_t j = 1; j < n; ++j) { + mixed ^= (uint64_t) ctx[j] * hp.ple_layer_multipliers[j]; + } + const int64_t base = (n - 2) * per_gram; + for (int64_t g = 0; g < per_gram; ++g) { + const int64_t h_i = base + g; + idx[i * n_heads + h_i] = + (int32_t) (mixed % hp.ple_head_vocab_sizes[h_i] + hp.ple_head_offsets[h_i]); + } + } + } + + ggml_backend_tensor_set(rows, idx.data(), 0, idx.size()*ggml_element_size(rows)); +} + +// Read a conv history out of its own recurrent row and write the new tail back. +// The shared build_conv_state cannot do this: qwen4exp has two such rows per layer. +ggml_tensor * llama_model_qwen4exp::graph::build_conv_state_at( + llm_graph_input_rs * inp, + ggml_tensor * conv_states_all, + ggml_tensor * x, + int64_t state_cols, + int64_t channels, + int il) { + const auto * mctx_cur = inp->mctx; + + const auto kv_head = mctx_cur->get_head(); + + const int64_t n_seqs = ubatch.n_seqs; + const int64_t row_total = conv_states_all->ne[0]; + + // the row is exactly this convolution's state, so the gather is reused as a whole + GGML_ASSERT(state_cols * channels == row_total); + + auto it = rs_rows.find(conv_states_all); + if (it == rs_rows.end()) { + it = rs_rows.emplace(conv_states_all, build_rs(inp, conv_states_all, row_total, n_seqs)).first; + } + ggml_tensor * rows = it->second; + + ggml_tensor * state = ggml_reshape_3d(ctx0, rows, state_cols, channels, n_seqs); + cb(state, "conv_state_at", il); + + ggml_tensor * conv_input = ggml_concat(ctx0, state, ggml_transpose(ctx0, x), 0); + + // keep the last state_cols columns for the next ubatch + const size_t row_size = ggml_row_size(conv_states_all->type, row_total); + + ggml_tensor * tail = ggml_view_3d(ctx0, conv_input, + state_cols, channels, n_seqs, + conv_input->nb[1], conv_input->nb[2], + ggml_row_size(conv_input->type, conv_input->ne[0] - state_cols)); + + ggml_tensor * dst = ggml_view_2d(ctx0, conv_states_all, + state_cols * channels, n_seqs, + conv_states_all->nb[1], + kv_head * row_size); + + ggml_build_forward_expand(gf, ggml_cpy(ctx0, ggml_cont(ctx0, tail), dst)); + + return conv_input; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_inp_ple( + const llama_memory_hybrid_idx_context * mctx_hyb) { + const int64_t n_heads = hparams.ple_n_heads; + + // the attention cells see every ubatch regardless of the layer types + auto ple_inp = std::make_unique( + static_cast(model), mctx_hyb->get_attn()); + + ple_inp->rows = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_heads * n_tokens); + ggml_set_input(ple_inp->rows); + ggml_tensor * rows = ple_inp->rows; + res->add_input(std::move(ple_inp)); + + // gather then flatten the heads: get_rows lays the head dimension out slowest, as the reference does + ggml_tensor * emb = ggml_get_rows(ctx0, model.per_layer_tok_embd, rows); + emb = ggml_reshape_2d(ctx0, emb, hparams.ple_head_dim * n_heads, n_tokens); + cb(emb, "ple_embd", -1); + + return emb; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_ple( + llm_graph_input_rs * inp, + ggml_tensor * emb, + ggml_tensor * hidden, + int il) { + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t hc_dim = hc * n_embd; + + ggml_tensor * key = build_lora_mm(model.layers[il].ple_key, emb); + ggml_tensor * value = build_lora_mm(model.layers[il].ple_value, emb); + + // both norms group over one hc stream, with a weight over the whole hc*n_embd layout + auto grouped_norm = [&](ggml_tensor * x, ggml_tensor * w) { + ggml_tensor * t = ggml_reshape_3d(ctx0, x, n_embd, hc, n_tokens); + t = ggml_rms_norm(ctx0, t, hparams.f_norm_rms_eps); + t = ggml_reshape_2d(ctx0, t, hc_dim, n_tokens); + t = ggml_mul(ctx0, t, w); + return ggml_reshape_3d(ctx0, t, n_embd, hc, n_tokens); + }; + + key = grouped_norm(key, model.layers[il].ple_norm_key); + ggml_tensor * query = grouped_norm(hidden, model.layers[il].ple_norm_query); + + // per-stream dot product, then a signed square root before the sigmoid + ggml_tensor * s = ggml_sum_rows(ctx0, ggml_mul(ctx0, key, query)); + s = ggml_scale(ctx0, s, 1.0f / sqrtf((float) n_embd)); + + ggml_tensor * mag = ggml_sqrt(ctx0, ggml_clamp(ctx0, ggml_abs(ctx0, s), 1e-6f, INFINITY)); + ggml_tensor * gate = ggml_sigmoid(ctx0, ggml_mul(ctx0, ggml_sgn(ctx0, s), mag)); + cb(gate, "ple_gate", il); + + // [n_embd, 1, T] value broadcast across the hc streams, scaled by the gate + ggml_tensor * v3 = ggml_reshape_3d(ctx0, value, n_embd, 1, n_tokens); + v3 = ggml_repeat_4d(ctx0, v3, n_embd, hc, n_tokens, 1); + + ggml_tensor * gated = ggml_mul(ctx0, v3, gate); + cb(gated, "ple_gated_value", il); + + ggml_tensor * normalized = grouped_norm( + ggml_reshape_2d(ctx0, gated, hc_dim, n_tokens), + model.layers[il].ple_norm_conv); + normalized = ggml_reshape_2d(ctx0, normalized, hc_dim, n_tokens); + + // depthwise causal conv, dilated by the n-gram size, as a sum of shifted copies + // ggml_conv_1d_dw is documented as unreliable: + // out[c, t] = sum_k w[k, c] * x[c, t - (K-1-k)*dilation] + // The history of the earlier ubatches is prepended, so a chunked prefill matches a single-shot one. + const int64_t kern = hparams.ple_conv_kernel; + const int64_t dil = hparams.ple_ngram_size; + const int64_t hist = (kern - 1) * dil; + + // the conv history is per sequence, so the input carries the sequence axis too + const int64_t n_seqs = ubatch.n_seqs; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + + // [hist + n_seq_tokens, hc_dim, n_seqs], tokens on ne[0] + ggml_tensor * padded = build_conv_state_at(inp, inp->mctx->get_p_l(il), + ggml_reshape_3d(ctx0, normalized, hc_dim, n_seq_tokens, n_seqs), + hist, hc_dim, il); + + ggml_tensor * conv_out = nullptr; + for (int64_t k = 0; k < kern; ++k) { + // tap k reads (kern-1-k)*dilation positions back + const int64_t start = hist - (kern - 1 - k) * dil; + + ggml_tensor * shifted = ggml_cont(ctx0, + ggml_transpose(ctx0, + ggml_view_3d(ctx0, padded, n_seq_tokens, hc_dim, n_seqs, + padded->nb[1], padded->nb[2], + ggml_row_size(padded->type, start)))); + + // column k of the [kern, hc_dim] kernel is one weight per channel + ggml_tensor * wk = ggml_cont(ctx0, + ggml_view_2d(ctx0, model.layers[il].ple_conv1d, 1, hc_dim, + model.layers[il].ple_conv1d->nb[1], + k * model.layers[il].ple_conv1d->nb[0])); + // this kernel keeps the file type, so cast it before it multiplies an f32 activation + wk = ggml_reshape_1d(ctx0, wk, hc_dim); + if (wk->type != GGML_TYPE_F32) { + wk = ggml_cast(ctx0, wk, GGML_TYPE_F32); + } + + ggml_tensor * term = ggml_mul(ctx0, shifted, wk); + conv_out = conv_out ? ggml_add(ctx0, conv_out, term) : term; + } + + conv_out = ggml_silu(ctx0, conv_out); + conv_out = ggml_reshape_3d(ctx0, ggml_cont(ctx0, conv_out), n_embd, hc, n_tokens); + cb(conv_out, "ple_conv_out", il); + + return ggml_add(ctx0, hidden, ggml_add(ctx0, gated, conv_out)); +} diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index b9f9d4b78af..fe3d14ffc55 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -149,6 +149,7 @@ if (LLAMA_LLGUIDANCE) endif () llama_build(test-recurrent-state-rollback.cpp) +llama_build(test-save-load-state.cpp) if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) # these tests are disabled on Windows because they use internal functions not exported with LLAMA_API (when building with shared libraries) @@ -237,6 +238,14 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) set_tests_properties(test-recurrent-state-rollback-dsv4 PROPERTIES FIXTURES_REQUIRED generate-models ) + + # Test state save/load functionality across all architectures, using the generated dummy models + llama_test( + test-save-load-state + LABEL main + ARGS --models "${MODEL_DIR}" + ) + set_tests_properties(test-save-load-state PROPERTIES FIXTURES_REQUIRED generate-models) endif() llama_build_and_test(test-chat-peg-parser.cpp peg-parser/simple-tokenize.cpp) @@ -299,10 +308,6 @@ llama_build_and_test(test-backend-sampler.cpp LABEL "model") llama_build_and_test(test-state-restore-fragmented.cpp LABEL "model" ARGS -m "${MODEL_DEST}") set_tests_properties(test-state-restore-fragmented PROPERTIES FIXTURES_REQUIRED test-download-model) -# Test state save/load functionality -llama_build_and_test(test-save-load-state.cpp LABEL "model" ARGS -m "${MODEL_DEST}") -set_tests_properties(test-save-load-state PROPERTIES FIXTURES_REQUIRED test-download-model) - if (APPLE) llama_build(test-rset-release.cpp) endif() diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index ba58f852eb4..e0907631abd 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -4,6 +4,7 @@ #include "llama.h" #include "speculative.h" +#include #include #include #include @@ -34,6 +35,62 @@ static void test(void) { std::numeric_limits::max(), std::numeric_limits::max()); + { + common_params_speculative spec; + spec.synth_len = 3.4; + + auto assert_invalid = [](const common_params_speculative & value, int32_t n_max) { + try { + common_speculative_synth_rates_resolve(&value, n_max); + assert(false); + } catch (const std::invalid_argument &) { + } + }; + + const auto rates = common_speculative_synth_rates_resolve(&spec, 4); + assert(rates.size() == 4); + assert(std::abs(rates[0] - 0.80581) < 1e-5); + assert(std::abs(rates[1] - 0.64933) < 1e-5); + assert(std::abs(rates[2] - 0.52323) < 1e-5); + assert(std::abs(rates[3] - 0.42163) < 1e-5); + assert(std::abs(1.0 + rates[0] + rates[1] + rates[2] + rates[3] - 3.4) < 1e-8); + + spec.synth_len = 1.0; + assert(common_speculative_synth_rates_resolve(&spec, 4) == std::vector({0.0, 0.0, 0.0, 0.0})); + + spec.synth_len = 5.0; + assert(common_speculative_synth_rates_resolve(&spec, 4) == std::vector({1.0, 1.0, 1.0, 1.0})); + + spec.synth_len = 5.1; + assert_invalid(spec, 4); + + spec.synth_len = std::numeric_limits::quiet_NaN(); + assert_invalid(spec, 4); + + spec.synth_len = 0.0; + assert_invalid(spec, 4); + + spec.synth_len = -1.0; + spec.synth_rates = {0.8, 0.6, 0.4}; + assert_invalid(spec, 4); + + spec.synth_rates = {0.8, 0.6, 0.4, 0.2}; + assert(common_speculative_synth_rates_resolve(&spec, 4) == spec.synth_rates); + + spec.synth_rates = {0.8, 0.9, 0.4, 0.2}; + assert_invalid(spec, 4); + + spec.synth_rates = {0.8, std::numeric_limits::quiet_NaN(), 0.4, 0.2}; + assert_invalid(spec, 4); + + spec.synth_rates = {0.8, 0.6, 0.4, -0.2}; + assert_invalid(spec, 4); + + spec.synth_rates = {0.8, 0.6, 0.4, 0.2}; + spec.synth_len = 3.0; + assert_invalid(spec, 4); + } + { common_params base; base.n_parallel = 4; @@ -197,6 +254,26 @@ static void test(void) { assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_SPECULATIVE)); assert(params.speculative.draft.n_max == 123); + { + common_params synth_params; + argv = {"binary_name", "--spec-synth-len", "3.4"}; + assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), synth_params, LLAMA_EXAMPLE_SERVER)); + assert(synth_params.speculative.synth_len == 3.4); + } + + { + common_params synth_params; + argv = {"binary_name", "--spec-synth-rates", "0.8,0.6,0.2"}; + assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), synth_params, LLAMA_EXAMPLE_SERVER)); + assert(synth_params.speculative.synth_rates == std::vector({0.8, 0.6, 0.2})); + } + + { + common_params synth_params; + argv = {"binary_name", "--spec-synth-len", "3.4x"}; + assert(false == common_params_parse(argv.size(), list_str_to_char(argv).data(), synth_params, LLAMA_EXAMPLE_SERVER)); + } + argv = {"binary_name", "-lm", "none"}; assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON)); assert(params.load_mode == LLAMA_LOAD_MODE_NONE); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 53e93a1448d..4a7a0623174 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -2469,8 +2469,13 @@ struct test_set_rows : public test_case { // See dicussion here: https://github.com/ggml-org/llama.cpp/pull/23760#issuecomment-4566312209 double max_nmse_err(ggml_backend_t backend) override { ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)); - if (type_dst == GGML_TYPE_Q8_0 && strcmp(ggml_backend_reg_name(reg), "WebGPU") == 0) { - return std::max(test_case::max_nmse_err(backend), 2e-7); + if (type_dst == GGML_TYPE_Q8_0) { + if (strcmp(ggml_backend_reg_name(reg), "WebGPU") == 0) { + return std::max(test_case::max_nmse_err(backend), 2e-7); + } + if (strcmp(ggml_backend_reg_name(reg), "HTP") == 0) { + return std::max(test_case::max_nmse_err(backend), 5e-6); + } } return test_case::max_nmse_err(backend); } @@ -8780,6 +8785,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_conv_transpose_2d({3, 2, 3, 1}, {2, 2, 1, 3}, 1, kernel_type)); test_cases.emplace_back(new test_conv_transpose_2d({10, 10, 9, 1}, {3, 3, 1, 9}, 2, kernel_type)); test_cases.emplace_back(new test_conv_transpose_2d({129, 63, 35, 1}, {3, 3, 48, 35}, 1, kernel_type)); + test_cases.emplace_back(new test_conv_transpose_2d({10, 10, 9, 2}, {3, 3, 1, 9}, 2, kernel_type)); // for multiple batches } test_cases.emplace_back(new test_count_equal(GGML_TYPE_F32, {4, 500, 1, 1})); @@ -9361,7 +9367,12 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_BF16, GGML_TYPE_F32, 16, 16, b, 50, 200, 64)); } - test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 1, 1, false, 8, 16, 1)); + // For issue 27873 + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_IQ2_XXS, GGML_TYPE_F32, 1, 1, false, 1, 8192, 4096)); + + for (int k : {1, 63, 65}) { + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_F16, GGML_TYPE_F32, 1, 1, false, 8, 16, k)); + } test_cases.emplace_back(new test_mul_mat_id_fusion(GGML_TYPE_F16, GGML_TYPE_F32, 16, 16, false, 32, 32, 32, 3)); // gpt-oss issue with Vulkan mmq_id @@ -9712,6 +9723,17 @@ static std::vector> make_test_cases_eval() { } } } + for (int k : {4, 8, 16, 32}) { + for (int nrows : {1, 8, 16}) { + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {202048, nrows, 1, 1}, k)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {151936, nrows, 1, 1}, k)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {8192, nrows, 1, 1}, k)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {8193, nrows, 1, 1}, k)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {8192, nrows, 1, 1}, k, true)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {202048, nrows, 1, 1}, k, true)); + } + } + for (int k : {1, 2, 3, 7, 15}) { test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {16, 10, 10, 10}, k)); test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {60, 10, 10, 10}, k)); @@ -10449,7 +10471,13 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_argsort(GGML_TYPE_F32, {200000, 16, 1, 1})); test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {2, 1, 1, 1}, 1)); - for (auto k : {1, 10, 40, 400}) { + // widths around the tiling threshold + for (auto cols : {4096, 8192, 12288, 16384, 24576, 32768, 65536, 131072}) { + for (auto nrows : {1, 16}) { + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, 16)); + } + } + for (auto k : {1, 4, 8, 10, 16, 32, 40, 400}) { for (auto nrows : {1, 16}) { for (auto cols : {k, 1000, 65000, 200000}) { test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, k)); diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index b8fd66ccae5..35a3286e4a1 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -65,7 +65,7 @@ static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) { } static void usage(char ** argv) { - printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-v/--verbose]\n", argv[0]); + printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-o/--out dir] [-v/--verbose] [-h/--help]\n", argv[0]); } static std::vector get_tokens(const uint32_t n_tokens, const uint32_t n_vocab, const size_t seed){ @@ -82,7 +82,7 @@ static std::vector get_tokens(const uint32_t n_tokens, const uint32 static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { gguf_context_ptr ret(gguf_init_empty()); llama_model_saver ms(arch, ret.get()); - const uint32_t n_ctx = 128; + const uint32_t n_ctx = 256; uint32_t n_vocab = 128; uint32_t n_embd = 256; @@ -249,8 +249,19 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { // MSA requires one indexer head per GQA (KV) head, unlike the DSA archs where the // indexer head count is independent of the main attention head count. - ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, arch == LLM_ARCH_MINIMAX_M3 || arch == LLM_ARCH_DEEPSEEK4 ? n_head : uint32_t(1)); - ms.add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, uint32_t(64)); + if (arch == LLM_ARCH_QWEN4EXP) { + ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT, uint32_t(4)); + ms.add_kv(LLM_KV_HYPER_CONNECTION_LOW_RANK, uint32_t(8)); + // without this the QSA layers fall back to dense and go uncovered + ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector(n_layer, 4)); + } + + // minimax-m3 keeps one indexer head per GQA head; the rest use a fixed 64 to match the fused + ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, arch == LLM_ARCH_MINIMAX_M3 ? n_head : uint32_t(64)); + // qwen4exp ropes indexer keys with the main rotary width, so its head can't be < n_rot + ms.add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, + arch == LLM_ARCH_QWEN4EXP ? n_embd_head : uint32_t(128)); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, uint32_t(8)); ms.add_kv(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, uint32_t(4)); ms.add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, uint32_t(1)); @@ -294,7 +305,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_XIELU_ALPHA_P, 1.0f); ms.add_kv(LLM_KV_XIELU_BETA, 1.0f); ms.add_kv(LLM_KV_XIELU_EPS, 1.0e-7f); - ms.add_kv(LLM_KV_SSM_INNER_SIZE, arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE ? 256 : 2*n_embd); + ms.add_kv(LLM_KV_SSM_INNER_SIZE, arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_QWEN4EXP ? 256 : 2*n_embd); ms.add_kv(LLM_KV_SSM_CONV_KERNEL, uint32_t(4)); ms.add_kv(LLM_KV_SSM_STATE_SIZE, uint32_t(128)); ms.add_kv(LLM_KV_SSM_TIME_STEP_RANK, n_head); @@ -411,6 +422,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_QWEN3NEXT: case LLM_ARCH_QWEN3VLMOE: case LLM_ARCH_QWEN35MOE: + case LLM_ARCH_QWEN4EXP: case LLM_ARCH_PHIMOE: case LLM_ARCH_DBRX: case LLM_ARCH_OLMOE: @@ -507,7 +519,7 @@ static bool arch_supported(const llm_arch arch) { } // FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI. #ifdef GGML_USE_WEBGPU - if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_DOTS3NOTE) { + if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_DOTS3NOTE || arch == LLM_ARCH_QWEN4EXP) { return false; } #endif // GGML_USE_WEBGPU @@ -752,6 +764,10 @@ int main(int argc, char ** argv) { std::string out; for (int i = 1; i < argc; i++) { + if (strcmp(argv[i], "-h") == 0 || strcmp(argv[i], "--help") == 0) { + usage(argv); + return 0; + } if (strcmp(argv[i], "-a") == 0 || strcmp(argv[i], "--arch") == 0) { if (i + 1 < argc) { const std::string arch_name = argv[++i]; diff --git a/tests/test-save-load-state.cpp b/tests/test-save-load-state.cpp index 6e93ce6fb8d..0ceab7c5452 100644 --- a/tests/test-save-load-state.cpp +++ b/tests/test-save-load-state.cpp @@ -3,8 +3,12 @@ #include "log.h" #include "llama-cpp.h" +#include #include +#include +#include #include +#include #include struct llama_batch_ptr { @@ -53,7 +57,9 @@ static llama_tokens generate_tokens(llama_context * ctx, llama_sampler * smpl, i // - decode the last token // - generate n_predict tokens static llama_tokens test_baseline(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens) { - auto ctx = llama_context_ptr{llama_init_from_model(model, common_context_params_to_llama(params))}; + auto params_ctx = common_context_params_to_llama(params); + params_ctx.n_seq_max = 2; + auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)}; auto sparams = llama_sampler_chain_default_params(); auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)}; @@ -161,7 +167,9 @@ static bool test_seq_rm_isolated( // - replay the last prompt token // - generate n_predict tokens and compare against expected result static bool test_state_load(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, const llama_tokens & expected_result) { - auto ctx = llama_context_ptr{llama_init_from_model(model, common_context_params_to_llama(params))}; + auto params_ctx = common_context_params_to_llama(params); + params_ctx.n_seq_max = 2; + auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)}; auto sparams = llama_sampler_chain_default_params(); auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)}; @@ -347,38 +355,18 @@ static bool test_seq_cp_device(struct llama_model * model, const struct common_p } -int main(int argc, char ** argv) { - std::setlocale(LC_NUMERIC, "C"); - - common_params params; - params.prompt = ""; - params.n_batch = 100; - params.out_file = "dump_state.bin"; - params.sampling.seed = 1234; - - common_init(); - - if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) { - return 1; - } - - if (params.n_parallel == 1) { - LOG_TRC("%s: n_parallel == 1, enabling unified kv cache\n", __func__); - params.kv_unified = true; - } - - if (params.n_predict < 0) { - params.n_predict = 16; - } - - ggml_backend_load_all(); +// Run the full save/load test suite (tests 1-5) for a single model. +// Returns true if all tests pass, false otherwise. +static bool run_save_load_tests_for_model(const std::string & model_path, const struct common_params & base_params) { + struct common_params params = base_params; + params.model.path = model_path; auto llama_init = common_init_from_params(params, true); auto * model = llama_init->model(); if (model == nullptr) { - LOG_ERR("%s: failed to init\n", __func__); - return 1; + LOG_ERR("%s: failed to init model '%s'\n", __func__, model_path.c_str()); + return false; } GGML_ASSERT(llama_init->context() == nullptr); @@ -411,30 +399,127 @@ int main(int argc, char ** argv) { // Test 1: baseline (saves state to disk) auto result_baseline = test_baseline(model, params, tokens); if (result_baseline.empty()) { - return 1; + return false; } // Test 2: sequence removal isolation if (!test_seq_rm_isolated(model, params, tokens)) { - return 1; + return false; } // Test 3: state load if (!test_state_load(model, params, tokens, result_baseline)) { - return 1; + return false; } // Test 4: seq copy (host) if (!test_seq_cp_host(model, params, tokens, result_baseline)) { - return 1; + return false; } // Test 5: seq copy (device) if (!test_seq_cp_device(model, params, tokens, result_baseline)) { - return 1; + return false; } LOG("\nAll tests passed.\n"); - return 0; + return true; +} + + +int main(int argc, char ** argv) { + std::setlocale(LC_NUMERIC, "C"); + + common_params params; + params.prompt = ""; + params.n_batch = 100; + params.out_file = "dump_state.bin"; + params.sampling.seed = 1234; + + common_init(); + + // extract our own --models DIR option before handing the rest to the common arg parser + std::string models_dir; + std::vector filtered_argv; + filtered_argv.push_back(argv[0]); + for (int i = 1; i < argc; i++) { + if (strcmp(argv[i], "--models") == 0) { + if (i + 1 >= argc) { + LOG_ERR("%s: --models requires a directory argument\n", __func__); + return 1; + } + models_dir = argv[i + 1]; + i++; + } else { + filtered_argv.push_back(argv[i]); + } + } + filtered_argv.push_back(nullptr); + const int fargc = (int)filtered_argv.size() - 1; + + // in --models mode there is no single model; set a placeholder so the common parser's + // "--model is required" check passes (each model is set individually inside the loop) + if (!models_dir.empty()) { + params.model.path = models_dir; + } + + if (!common_params_parse(fargc, filtered_argv.data(), params, LLAMA_EXAMPLE_COMMON)) { + return 1; + } + + if (params.n_parallel == 1) { + LOG_TRC("%s: n_parallel == 1, enabling unified kv cache\n", __func__); + params.kv_unified = true; + } + + if (params.n_predict < 0) { + params.n_predict = 16; + } + + ggml_backend_load_all(); + + if (!models_dir.empty()) { + // run the suite over every dummy model in the directory + if (!std::filesystem::exists(models_dir) || !std::filesystem::is_directory(models_dir)) { + LOG_ERR("%s: models directory '%s' does not exist\n", __func__, models_dir.c_str()); + return 1; + } + + std::vector models; + for (const auto & entry : std::filesystem::directory_iterator(models_dir)) { + if (entry.is_regular_file() && entry.path().extension() == ".gguf") { + models.push_back(entry.path().string()); + } + } + std::sort(models.begin(), models.end()); + + if (models.empty()) { + LOG_ERR("%s: no .gguf models found in '%s'\n", __func__, models_dir.c_str()); + return 1; + } + + LOG_INF("%s: running save/load tests over %zu models in '%s'\n", __func__, models.size(), models_dir.c_str()); + + size_t n_pass = 0; + size_t n_fail = 0; + for (const auto & model_path : models) { + LOG("\n================================================================\n"); + LOG_INF("%s: model %s\n", __func__, model_path.c_str()); + + if (run_save_load_tests_for_model(model_path, params)) { + n_pass++; + } else { + n_fail++; + } + } + + LOG("\n================================================================\n"); + LOG_INF("%s: summary: %zu passed, %zu failed (of %zu)\n", __func__, n_pass, n_fail, models.size()); + + return n_fail == 0 ? 0 : 1; + } + + // single-model mode + return run_save_load_tests_for_model(params.model.path, params) ? 0 : 1; } diff --git a/tools/cli/README.md b/tools/cli/README.md index c9cbacafcd1..163ee4fbaf0 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -59,12 +59,14 @@ | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)
(env: LLAMA_ARG_MMAP) | | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available
(env: LLAMA_ARG_DIO) | | `-lm, --load-mode MODE` | model loading mode (default: auto)
- auto: mmap, unless a device does not support it
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | +| `--tensor-read-lazy MODE` | on-demand reading of certain tensors, for example per-layer embeddings (default: auto)
- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)
- auto: on, but only for tensors larger than 4 GiB
- off: always keep them resident
(env: LLAMA_ARG_TENSOR_READ_LAZY) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | | `-ot, --override-tensor =,...` | override tensor buffer type
(env: LLAMA_ARG_OVERRIDE_TENSOR) | | `-cmoe, --cpu-moe` | keep all Mixture of Experts (MoE) weights in the CPU
(env: LLAMA_ARG_CPU_MOE) | | `-ncmoe, --n-cpu-moe N` | keep the Mixture of Experts (MoE) weights of the first N layers in the CPU
(env: LLAMA_ARG_N_CPU_MOE) | +| `-ncffn, --n-cpu-ffn N` | keep the dense FFN weights of the first N layers in the CPU
(dense models; for MoE expert weights use --n-cpu-moe)
(env: LLAMA_ARG_N_CPU_FFN) | | `-ngl, --gpu-layers, --n-gpu-layers N` | max. number of layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)
(env: LLAMA_ARG_N_GPU_LAYERS) | | `-sm, --split-mode {none,layer,row,tensor}` | how to split the model across multiple GPUs, one of:
- none: use one GPU only
- layer (default): split layers and KV across GPUs (pipelined)
- row: split weight across GPUs by rows (parallelized)
- tensor: split weights and KV across GPUs (parallelized, EXPERIMENTAL)
(env: LLAMA_ARG_SPLIT_MODE) | | `-ts, --tensor-split N0,N1,N2,...` | fraction of the model to offload to each GPU, comma-separated list of proportions, e.g. 3,1
(env: LLAMA_ARG_TENSOR_SPLIT) | @@ -154,7 +156,6 @@ | `-sysf, --system-prompt-file FNAME` | a file containing the system prompt (default: none) | | `-r, --reverse-prompt PROMPT` | halt generation at PROMPT, return control in interactive mode | | `-sp, --special` | special tokens output enabled (default: false) | -| `-cnv, --conversation, -no-cnv, --no-conversation` | whether to run in conversation mode:
- does not print special tokens and suffix/prefix
- interactive mode is also enabled
(default: auto enabled if chat template is available) | | `-st, --single-turn` | run conversation for a single turn only, then exit when done
will not be interactive if first turn is predefined with --prompt
(default: false) | | `-mli, --multiline-input` | allows you to write or paste multiple lines without ending each in '\' | | `--warmup, --no-warmup` | whether to perform warmup with an empty run (default: enabled) | @@ -166,6 +167,9 @@ | `--image, --audio, --video FILE` | path to an image, audio, or video file. use with multimodal models, use comma-separated values for multiple files | | `--image-min-tokens N` | minimum number of tokens each image can take, only used by vision models with dynamic resolution (default: read from model)
(env: LLAMA_ARG_IMAGE_MIN_TOKENS) | | `--image-max-tokens N` | maximum number of tokens each image can take, only used by vision models with dynamic resolution (default: read from model)
(env: LLAMA_ARG_IMAGE_MAX_TOKENS) | +| `--video-fps N` | target video frame rate (default: 4.0)
(env: LLAMA_ARG_VIDEO_FPS) | +| `--video-timestamp-interval N` | interval in milliseconds between text timestamps (default: 5000)
(env: LLAMA_ARG_VIDEO_TIMESTAMP_INTERVAL) | +| `--video-ffmpeg-dir DIR` | path to the directory containing ffmpeg and ffprobe (default: search in PATH)
(env: LLAMA_ARG_VIDEO_FFMPEG_DIR) | | `-o, --output, --output-file FNAME` | output file (default: '') | | `--chat-template-kwargs STRING` | sets additional params for the json template parser, must be a valid json object string, e.g. '{"key1":"value1","key2":"value2"}'
(env: LLAMA_ARG_CHAT_TEMPLATE_KWARGS) | | `--jinja, --no-jinja` | whether to use jinja template engine for chat (default: enabled)
(env: LLAMA_ARG_JINJA) | @@ -197,6 +201,8 @@ | `--spec-draft-n-cpu-moe, --spec-draft-ncmoe, -ncmoed, --n-cpu-moe-draft N` | keep the Mixture of Experts (MoE) weights of the first N layers in the CPU for the draft model
(env: LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE) | | `--spec-draft-n-max N` | number of tokens to draft for speculative decoding (default: 3)
(env: LLAMA_ARG_SPEC_DRAFT_N_MAX) | | `--spec-draft-n-min N` | minimum number of draft tokens to use for speculative decoding (default: 0)
(env: LLAMA_ARG_SPEC_DRAFT_N_MIN) | +| `--spec-synth-len L` | target mean synthetic acceptance length, including the target token (benchmarking only)
(env: LLAMA_ARG_SPEC_SYNTH_LEN) | +| `--spec-synth-rates P0,P1,...` | comma-separated unconditional per-position synthetic acceptance probabilities (benchmarking only)
(env: LLAMA_ARG_SPEC_SYNTH_RATES) | | `--spec-draft-p-split, --draft-p-split P` | speculative decoding split probability (default: 0.10)
(env: LLAMA_ARG_SPEC_DRAFT_P_SPLIT) | | `--spec-draft-p-min, --draft-p-min P` | minimum speculative decoding probability (greedy) (default: 0.00)
(env: LLAMA_ARG_SPEC_DRAFT_P_MIN) | | `--spec-draft-backend-sampling, --no-spec-draft-backend-sampling` | offload draft sampling to the backend (default: enabled)
(env: LLAMA_ARG_SPEC_DRAFT_BACKEND_SAMPLING) | diff --git a/tools/completion/README.md b/tools/completion/README.md index 833687dcad4..0cd86bac70f 100644 --- a/tools/completion/README.md +++ b/tools/completion/README.md @@ -142,12 +142,14 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1 | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)
(env: LLAMA_ARG_MMAP) | | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available
(env: LLAMA_ARG_DIO) | | `-lm, --load-mode MODE` | model loading mode (default: auto)
- auto: mmap, unless a device does not support it
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | +| `--tensor-read-lazy MODE` | on-demand reading of certain tensors, for example per-layer embeddings (default: auto)
- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)
- auto: on, but only for tensors larger than 4 GiB
- off: always keep them resident
(env: LLAMA_ARG_TENSOR_READ_LAZY) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | | `-ot, --override-tensor =,...` | override tensor buffer type
(env: LLAMA_ARG_OVERRIDE_TENSOR) | | `-cmoe, --cpu-moe` | keep all Mixture of Experts (MoE) weights in the CPU
(env: LLAMA_ARG_CPU_MOE) | | `-ncmoe, --n-cpu-moe N` | keep the Mixture of Experts (MoE) weights of the first N layers in the CPU
(env: LLAMA_ARG_N_CPU_MOE) | +| `-ncffn, --n-cpu-ffn N` | keep the dense FFN weights of the first N layers in the CPU
(dense models; for MoE expert weights use --n-cpu-moe)
(env: LLAMA_ARG_N_CPU_FFN) | | `-ngl, --gpu-layers, --n-gpu-layers N` | max. number of layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)
(env: LLAMA_ARG_N_GPU_LAYERS) | | `-sm, --split-mode {none,layer,row,tensor}` | how to split the model across multiple GPUs, one of:
- none: use one GPU only
- layer (default): split layers and KV across GPUs (pipelined)
- row: split weight across GPUs by rows (parallelized)
- tensor: split weights and KV across GPUs (parallelized, EXPERIMENTAL)
(env: LLAMA_ARG_SPLIT_MODE) | | `-ts, --tensor-split N0,N1,N2,...` | fraction of the model to offload to each GPU, comma-separated list of proportions, e.g. 3,1
(env: LLAMA_ARG_TENSOR_SPLIT) | diff --git a/tools/llama-bench/README.md b/tools/llama-bench/README.md index 42cb14859f0..a1404d2e33e 100644 --- a/tools/llama-bench/README.md +++ b/tools/llama-bench/README.md @@ -67,6 +67,7 @@ test parameters: -nkvo, --no-kv-offload <0|1> (default: 0) -fa, --flash-attn (default: auto) -dev, --device (default: auto) + --tensor-read-lazy (default: auto) -mmp, --mmap <0|1> (DEPRECATED IN FAVOUR OF --load-mode) -dio, --direct-io <0|1> (DEPRECATED IN FAVOUR OF --load-mode) -embd, --embeddings <0|1> (default: 0) diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index a2da93b9a28..1b4bbde4d4f 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -271,6 +271,19 @@ static const char * split_mode_str(llama_split_mode mode) { } } +static const char * lazy_mode_str(llama_lazy_mode mode) { + switch (mode) { + case LLAMA_LAZY_MODE_OFF: + return "off"; + case LLAMA_LAZY_MODE_AUTO: + return "auto"; + case LLAMA_LAZY_MODE_ON: + return "on"; + default: + GGML_ABORT("invalid tensor read lazy mode"); + } +} + static std::string pair_str(const std::pair & p) { static char buf[32]; snprintf(buf, sizeof(buf), "%d,%d", p.first, p.second); @@ -341,6 +354,7 @@ struct cmd_params { std::vector n_cpu_moe; std::vector split_mode; std::vector load_mode; + std::vector lazy_mode; std::vector main_gpu; std::vector no_kv_offload; std::vector flash_attn; @@ -385,6 +399,7 @@ static const cmd_params cmd_params_defaults = { /* n_cpu_moe */ { 0 }, /* split_mode */ { LLAMA_SPLIT_MODE_LAYER }, /* load_mode */ { LLAMA_LOAD_MODE_AUTO }, + /* lazy_mode */ { LLAMA_LAZY_MODE_AUTO }, /* main_gpu */ { 0 }, /* no_kv_offload */ { false }, /* flash_attn */ { LLAMA_FLASH_ATTN_TYPE_AUTO }, @@ -460,6 +475,7 @@ static void print_usage(int /* argc */, char ** argv) { printf(" -fa, --flash-attn (default: %s)\n", join(transform_to_str(cmd_params_defaults.flash_attn, llama_flash_attn_type_name), ",").c_str()); printf(" -dev, --device (default: auto)\n"); printf(" -lm, --load-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str()); + printf(" --tensor-read-lazy (default: %s)\n", join(transform_to_str(cmd_params_defaults.lazy_mode, lazy_mode_str), ",").c_str()); printf(" -mmp, --mmap <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n"); printf(" -dio, --direct-io <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n"); printf(" -embd, --embeddings <0|1> (default: %s)\n", join(cmd_params_defaults.embeddings, ",").c_str()); @@ -786,6 +802,32 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { break; } params.load_mode.insert(params.load_mode.end(), modes.begin(), modes.end()); + } else if (arg == "--tensor-read-lazy") { + if (++i >= argc) { + invalid_param = true; + break; + } + auto p = string_split(argv[i], split_delim); + + std::vector modes; + for (const auto & m : p) { + llama_lazy_mode mode; + if (m == "on") { + mode = LLAMA_LAZY_MODE_ON; + } else if (m == "auto") { + mode = LLAMA_LAZY_MODE_AUTO; + } else if (m == "off") { + mode = LLAMA_LAZY_MODE_OFF; + } else { + invalid_param = true; + break; + } + modes.push_back(mode); + } + if (invalid_param) { + break; + } + params.lazy_mode.insert(params.lazy_mode.end(), modes.begin(), modes.end()); } else if (arg == "-mg" || arg == "--main-gpu") { if (++i >= argc) { invalid_param = true; @@ -1137,6 +1179,9 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { if (params.load_mode.empty()) { params.load_mode = cmd_params_defaults.load_mode; } + if (params.lazy_mode.empty()) { + params.lazy_mode = cmd_params_defaults.lazy_mode; + } if (params.main_gpu.empty()) { params.main_gpu = cmd_params_defaults.main_gpu; } @@ -1203,6 +1248,7 @@ struct cmd_params_instance { int n_cpu_moe; llama_split_mode split_mode; llama_load_mode load_mode; + llama_lazy_mode lazy_mode; int main_gpu; bool no_kv_offload; llama_flash_attn_type flash_attn; @@ -1224,6 +1270,7 @@ struct cmd_params_instance { } mparams.split_mode = split_mode; mparams.load_mode = load_mode; + mparams.lazy_mode = lazy_mode; mparams.main_gpu = main_gpu; mparams.tensor_split = tensor_split.data(); mparams.no_host = no_host; @@ -1254,7 +1301,7 @@ struct cmd_params_instance { merged.reserve(merged.size() + (size_t) n_cpu_moe + 1); for (int i = 0; i < n_cpu_moe; ++i) { - patterns.push_back(llm_ffn_exps_block_regex(i)); + patterns.push_back(llm_ffn_block_regex(i, LLM_FFN_EXPS_REGEX)); merged.push_back({ patterns.back().c_str(), ggml_backend_cpu_buffer_type() }); } @@ -1271,7 +1318,8 @@ struct cmd_params_instance { return model == other.model && n_gpu_layers == other.n_gpu_layers && n_cpu_moe == other.n_cpu_moe && split_mode == other.split_mode && main_gpu == other.main_gpu && tensor_split == other.tensor_split && - load_mode == other.load_mode && devices == other.devices && no_host == other.no_host && + load_mode == other.load_mode && lazy_mode == other.lazy_mode && + devices == other.devices && no_host == other.no_host && vec_tensor_buft_override_equal(tensor_buft_overrides, other.tensor_buft_overrides); } @@ -1305,6 +1353,7 @@ static std::vector get_cmd_params_instances(const cmd_param for (const auto & ncmoe : params.n_cpu_moe) for (const auto & sm : params.split_mode) for (const auto & lm : params.load_mode) + for (const auto & lzm : params.lazy_mode) for (const auto & mg : params.main_gpu) for (const auto & devs : params.devices) for (const auto & ts : params.tensor_split) @@ -1344,6 +1393,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .n_cpu_moe = */ ncmoe, /* .split_mode = */ sm, /* .load_mode = */ lm, + /* .lazy_mode = */ lzm, /* .main_gpu = */ mg, /* .no_kv_offload = */ nkvo, /* .flash_attn = */ fa, @@ -1380,6 +1430,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .n_cpu_moe = */ ncmoe, /* .split_mode = */ sm, /* .load_mode = */ lm, + /* .lazy_mode = */ lzm, /* .main_gpu = */ mg, /* .no_kv_offload = */ nkvo, /* .flash_attn = */ fa, @@ -1416,6 +1467,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .n_cpu_moe = */ ncmoe, /* .split_mode = */ sm, /* .load_mode = */ lm, + /* .lazy_mode = */ lzm, /* .main_gpu = */ mg, /* .no_kv_offload = */ nkvo, /* .flash_attn = */ fa, @@ -1457,6 +1509,7 @@ struct test { int n_cpu_moe; llama_split_mode split_mode; llama_load_mode load_mode; + llama_lazy_mode lazy_mode; int main_gpu; bool no_kv_offload; llama_flash_attn_type flash_attn; @@ -1496,6 +1549,7 @@ struct test { n_cpu_moe = inst.n_cpu_moe; split_mode = inst.split_mode; load_mode = inst.load_mode; + lazy_mode = inst.lazy_mode; main_gpu = inst.main_gpu; no_kv_offload = inst.no_kv_offload; flash_attn = inst.flash_attn; @@ -1563,7 +1617,8 @@ struct test { "n_ubatch", "n_threads", "cpu_mask", "cpu_strict", "poll", "type_k", "type_v", "n_gpu_layers", "n_cpu_moe", "split_mode", "main_gpu", "no_kv_offload", "flash_attn", "devices", "tensor_split", - "tensor_buft_overrides", "load_mode", "embeddings", + "tensor_buft_overrides", "load_mode", "lazy_mode", + "embeddings", "no_op_offload", "no_host", "fit_target", "fit_min_ctx", "n_prompt", "n_gen", "n_depth", "test_time", "avg_ns", "stddev_ns", "avg_ts", "stddev_ts" @@ -1588,7 +1643,7 @@ struct test { if (field == "avg_ts" || field == "stddev_ts") { return FLOAT; } - if (field == "load_mode") { + if (field == "load_mode" || field == "lazy_mode") { return STRING; } return STRING; @@ -1658,6 +1713,7 @@ struct test { tensor_split_str, tensor_buft_overrides_str, llama_load_mode_name(load_mode), + lazy_mode_str(lazy_mode), std::to_string(embeddings), std::to_string(no_op_offload), std::to_string(no_host), @@ -1972,6 +2028,9 @@ struct markdown_printer : public printer { if (params.load_mode.size() > 1 || params.load_mode != cmd_params_defaults.load_mode) { fields.emplace_back("load_mode"); } + if (params.lazy_mode.size() > 1 || params.lazy_mode != cmd_params_defaults.lazy_mode) { + fields.emplace_back("lazy_mode"); + } if (params.embeddings.size() > 1 || params.embeddings != cmd_params_defaults.embeddings) { fields.emplace_back("embeddings"); } diff --git a/tools/mtmd/mtmd-cli.cpp b/tools/mtmd/mtmd-cli.cpp index f6c787fdb68..97678c6b2e5 100644 --- a/tools/mtmd/mtmd-cli.cpp +++ b/tools/mtmd/mtmd-cli.cpp @@ -87,6 +87,9 @@ struct mtmd_cli_context { mtmd::bitmaps bitmaps; std::vector videos; + mtmd_helper_init_opt init_opt = mtmd_helper_init_opt_default(); + std::string video_ffmpeg_bin_dir; + mtmd::batch_ptr mbatch; // chat template @@ -170,6 +173,12 @@ struct mtmd_cli_context { LOG_ERR("Failed to load vision model from %s\n", clip_path); exit(1); } + + video_ffmpeg_bin_dir = params.video_ffmpeg_bin_dir; + init_opt.video_params.fps_target = params.video_fps; + init_opt.video_params.timestamp_interval_ms = params.video_timestamp_interval_ms; + init_opt.video_params.ffmpeg_bin_dir = video_ffmpeg_bin_dir.empty() + ? nullptr : video_ffmpeg_bin_dir.c_str(); } bool check_antiprompt(const llama_tokens & generated_tokens) { @@ -184,7 +193,7 @@ struct mtmd_cli_context { } bool load_media(const std::string & fname) { - auto res = mtmd_helper_bitmap_init_from_file(ctx_vision.get(), fname.c_str(), false); + auto res = mtmd_helper_bitmap_init_from_file(ctx_vision.get(), fname.c_str(), false, init_opt); if (!res.bitmap) { return false; } diff --git a/tools/mtmd/mtmd-helper.cpp b/tools/mtmd/mtmd-helper.cpp index f1defb64772..77f9d58fe23 100644 --- a/tools/mtmd/mtmd-helper.cpp +++ b/tools/mtmd/mtmd-helper.cpp @@ -369,14 +369,18 @@ static bool is_webp_file(const unsigned char * buf, size_t len) { } #ifdef MTMD_VIDEO -static mtmd_bitmap * decode_webp_with_ffmpeg(mtmd_context * mctx, const unsigned char * buf, size_t len, bool placeholder); +static mtmd_bitmap * decode_webp_with_ffmpeg(mtmd_context * mctx, const unsigned char * buf, size_t len, bool placeholder, + const mtmd_helper_video_init_params & params); #endif -mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, const unsigned char * buf, size_t len, bool placeholder) { +mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, const unsigned char * buf, size_t len, bool placeholder, + mtmd_helper_init_opt opt) { // calculate the hash if needed std::string id; mtmd_bitmap * result = nullptr; + GGML_UNUSED(opt); // only used by video code paths + if (!placeholder) { // use sha256 to prevent cache poisoning id = hash_sha256_hex(buf, len); @@ -414,7 +418,7 @@ mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, #ifdef MTMD_VIDEO // stb_image does not support webp; decode it with ffmpeg as a single frame if (!result && is_webp_file(buf, len)) { - result = decode_webp_with_ffmpeg(ctx, buf, len, placeholder); + result = decode_webp_with_ffmpeg(ctx, buf, len, placeholder, opt.video_params); if (!result) { LOG_ERR("%s: failed to decode webp buffer\n", __func__); return {nullptr, nullptr}; @@ -427,8 +431,7 @@ mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, // last try: load as video #ifdef MTMD_VIDEO if (!result) { - auto params = mtmd_helper_video_init_params_default(); - auto video_ctx = mtmd_helper_video_init_from_buf(ctx, buf, len, params); + auto video_ctx = mtmd_helper_video_init_from_buf(ctx, buf, len, opt.video_params); if (!video_ctx) { LOG_ERR("%s: failed to decode buffer as either image/audio/video\n", __func__); return {nullptr, nullptr}; @@ -456,7 +459,8 @@ mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, return {nullptr, nullptr}; } -mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_file(mtmd_context * ctx, const char * fname, bool placeholder) { +mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_file(mtmd_context * ctx, const char * fname, bool placeholder, + mtmd_helper_init_opt opt) { #ifdef _WIN32 int wlen = MultiByteToWideChar(CP_UTF8, 0, fname, -1, NULL, 0); if (!wlen) { @@ -497,7 +501,7 @@ mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_file(mtmd_context * ctx, return {nullptr, nullptr}; } - return mtmd_helper_bitmap_init_from_buf(ctx, buf.data(), buf.size(), placeholder); + return mtmd_helper_bitmap_init_from_buf(ctx, buf.data(), buf.size(), placeholder, opt); } bool mtmd_helper_support_video(mtmd_context * ctx) { @@ -855,6 +859,12 @@ mtmd_helper_video_init_params mtmd_helper_video_init_params_default() { }; } +mtmd_helper_init_opt mtmd_helper_init_opt_default() { + return { + /* video_params */ mtmd_helper_video_init_params_default(), + }; +} + static std::string video_resolve_bin(const char * bin_dir, const char * name) { if (!bin_dir || bin_dir[0] == '\0') { return name; // rely on PATH @@ -876,8 +886,8 @@ static std::string video_resolve_bin(const char * bin_dir, const char * name) { } #ifdef MTMD_VIDEO -static mtmd_bitmap * decode_webp_with_ffmpeg(mtmd_context * mctx, const unsigned char * buf, size_t len, bool placeholder) { - auto params = mtmd_helper_video_init_params_default(); +static mtmd_bitmap * decode_webp_with_ffmpeg(mtmd_context * mctx, const unsigned char * buf, size_t len, bool placeholder, + const mtmd_helper_video_init_params & params) { mtmd_helper_video vctx; vctx.mctx = mctx; vctx.input_buf.assign(buf, buf + len); diff --git a/tools/mtmd/mtmd-helper.h b/tools/mtmd/mtmd-helper.h index 58dfb152501..772e0f091bd 100644 --- a/tools/mtmd/mtmd-helper.h +++ b/tools/mtmd/mtmd-helper.h @@ -23,6 +23,23 @@ extern "C" { struct mtmd_helper_video; typedef struct mtmd_helper_video mtmd_helper_video; +struct mtmd_helper_video_init_params { + float fps_target; // desired output fps; <= 0 means use the video's native fps, defaulted to 4.0f + const char * ffmpeg_bin_dir; // directory containing ffmpeg/ffprobe binaries; NULL means search PATH + int64_t timestamp_interval_ms; // interval for adding timestamp as text chunk (example: "[10m50.5s]"); <= 0 means no timestamp, defaulted to 5000ms + // TODO @ngxson : allow "placeholder" bitmap output for counting tokens +}; + +MTMD_API struct mtmd_helper_video_init_params mtmd_helper_video_init_params_default(void); + +// opt for mtmd_helper_bitmap_init_from_*() +struct mtmd_helper_init_opt { + struct mtmd_helper_video_init_params video_params; +}; +typedef struct mtmd_helper_init_opt mtmd_helper_init_opt; + +MTMD_API struct mtmd_helper_init_opt mtmd_helper_init_opt_default(void); + // Set callback for all future logging events. // If this is not called, or NULL is supplied, everything is output on stderr. // Note: this also call mtmd_log_set() internally @@ -40,7 +57,11 @@ struct mtmd_helper_bitmap_wrapper { // it calls mtmd_helper_bitmap_init_from_buf() internally // returns nullptr on failure // this function is thread-safe -MTMD_API struct mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_file(mtmd_context * ctx, const char * fname, bool placeholder); +MTMD_API struct mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_file( + mtmd_context * ctx, + const char * fname, + bool placeholder, + struct mtmd_helper_init_opt opt); // helper function to construct a mtmd_bitmap from a buffer containing a file // supported formats: @@ -53,7 +74,11 @@ MTMD_API struct mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_file(mtm // - output bitmap will have SHA-256 hash (hex string) as the ID // returns nullptr on failure // this function is thread-safe -MTMD_API struct mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, const unsigned char * buf, size_t len, bool placeholder); +MTMD_API struct mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf( + mtmd_context * ctx, + const unsigned char * buf, size_t len, + bool placeholder, + struct mtmd_helper_init_opt opt); // helper to count the total number of tokens from a list of chunks, useful to keep track of KV cache MTMD_API size_t mtmd_helper_get_n_tokens(const mtmd_input_chunks * chunks); @@ -124,14 +149,7 @@ struct mtmd_helper_video_info { int32_t n_frames; // estimated total frames at effective fps (-1 if unknown) }; -struct mtmd_helper_video_init_params { - float fps_target; // desired output fps; <= 0 means use the video's native fps, defaulted to 4.0f - const char * ffmpeg_bin_dir; // directory containing ffmpeg/ffprobe binaries; NULL means search PATH - int64_t timestamp_interval_ms; // interval for adding timestamp as text chunk (example: "[10m50.5s]"); <= 0 means no timestamp, defaulted to 5000ms - // TODO @ngxson : allow "placeholder" bitmap output for counting tokens -}; - -MTMD_API struct mtmd_helper_video_init_params mtmd_helper_video_init_params_default(void); +// note: mtmd_helper_video_init_params is defined at the top, as it is part of mtmd_helper_init_opt // returns NULL on failure (ffprobe not found, file unreadable, etc.) MTMD_API mtmd_helper_video * mtmd_helper_video_init( diff --git a/tools/quantize/quantize.cpp b/tools/quantize/quantize.cpp index 8d03c8fcd42..38950036cd8 100644 --- a/tools/quantize/quantize.cpp +++ b/tools/quantize/quantize.cpp @@ -122,7 +122,7 @@ static bool try_parse_ftype(const std::string & ftype_str_in, llama_ftype & ftyp static void usage(const char * executable) { printf("usage: %s [--help] [--allow-requantize] [--leave-output-tensor] [--pure] [--imatrix] [--include-weights]\n", executable); printf(" [--exclude-weights] [--output-tensor-type] [--token-embedding-type] [--tensor-type] [--tensor-type-file]\n"); - printf(" [--prune-layers] [--keep-split] [--override-kv] [--dry-run]\n"); + printf(" [--prune-layers] [--keep-split] [--override-kv] [--dry-run] [--max-buffer-size]\n"); printf(" model-f32.gguf [model-quant.gguf] type [nthreads]\n\n"); printf(" --allow-requantize\n"); printf(" allow requantizing tensors that have already been quantized\n"); @@ -161,7 +161,10 @@ static void usage(const char * executable) { printf(" WARNING: this is an advanced option, use with care.\n"); printf(" --dry-run\n"); printf(" calculate and show the final quantization size without performing quantization\n"); - printf(" example: llama-quantize --dry-run model-f32.gguf Q4_K\n\n"); + printf(" example: llama-quantize --dry-run model-f32.gguf Q4_K\n"); + printf(" --max-buffer-size MiB\n"); + printf(" max amount of tensor rows kept in memory while quantizing one tensor (default: 8192)\n"); + printf(" lower it to quantize models with very large tensors on a machine with little RAM\n\n"); printf("note: --include-weights and --exclude-weights cannot be used together\n\n"); printf("-----------------------------------------------------------------------------\n"); printf(" allowed quantization types\n"); @@ -467,6 +470,16 @@ int llama_quantize(int argc, char ** argv) { } } else if (strcmp(argv[arg_idx], "--keep-split") == 0) { params.keep_split = true; + } else if (strcmp(argv[arg_idx], "--max-buffer-size") == 0) { + if (arg_idx == argc-1) { + usage(argv[0]); + } + const int mib = atoi(argv[++arg_idx]); + if (mib <= 0) { + fprintf(stderr, "%s: invalid --max-buffer-size '%s'\n", __func__, argv[arg_idx]); + return 1; + } + params.max_buf_size = (size_t) mib * 1024 * 1024; } else { usage(argv[0]); } diff --git a/tools/server/README.md b/tools/server/README.md index 93736c3edfa..3c2228f3432 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -76,12 +76,14 @@ For the full list of features, please refer to [server's changelog](https://gith | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)
(env: LLAMA_ARG_MMAP) | | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available
(env: LLAMA_ARG_DIO) | | `-lm, --load-mode MODE` | model loading mode (default: auto)
- auto: mmap, unless a device does not support it
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | +| `--tensor-read-lazy MODE` | on-demand reading of certain tensors, for example per-layer embeddings (default: auto)
- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)
- auto: on, but only for tensors larger than 4 GiB
- off: always keep them resident
(env: LLAMA_ARG_TENSOR_READ_LAZY) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | | `-ot, --override-tensor =,...` | override tensor buffer type
(env: LLAMA_ARG_OVERRIDE_TENSOR) | | `-cmoe, --cpu-moe` | keep all Mixture of Experts (MoE) weights in the CPU
(env: LLAMA_ARG_CPU_MOE) | | `-ncmoe, --n-cpu-moe N` | keep the Mixture of Experts (MoE) weights of the first N layers in the CPU
(env: LLAMA_ARG_N_CPU_MOE) | +| `-ncffn, --n-cpu-ffn N` | keep the dense FFN weights of the first N layers in the CPU
(dense models; for MoE expert weights use --n-cpu-moe)
(env: LLAMA_ARG_N_CPU_FFN) | | `-ngl, --gpu-layers, --n-gpu-layers N` | max. number of layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)
(env: LLAMA_ARG_N_GPU_LAYERS) | | `-sm, --split-mode {none,layer,row,tensor}` | how to split the model across multiple GPUs, one of:
- none: use one GPU only
- layer (default): split layers and KV across GPUs (pipelined)
- row: split weight across GPUs by rows (parallelized)
- tensor: split weights and KV across GPUs (parallelized, EXPERIMENTAL)
(env: LLAMA_ARG_SPLIT_MODE) | | `-ts, --tensor-split N0,N1,N2,...` | fraction of the model to offload to each GPU, comma-separated list of proportions, e.g. 3,1
(env: LLAMA_ARG_TENSOR_SPLIT) | @@ -161,6 +163,7 @@ For the full list of features, please refer to [server's changelog](https://gith | -------- | ----------- | | `-lcs, --lookup-cache-static FNAME` | path to static lookup cache to use for lookup decoding (not updated by generation) | | `-lcd, --lookup-cache-dynamic FNAME` | path to dynamic lookup cache to use for lookup decoding (updated by generation) | +| `--kv-unified-per-slot N` | context limit per parallel slot (default: unset, behavior unchanged).
when set without -c/--ctx-size, the shared KV pool is sized to n_parallel*N
(env: LLAMA_ARG_KV_UNIFIED_PER_SLOT) | | `-ctxcp, --ctx-checkpoints, --swa-checkpoints N` | max number of context checkpoints to create per slot (default: 32)[(more info)](https://github.com/ggml-org/llama.cpp/pull/15293)
(env: LLAMA_ARG_CTX_CHECKPOINTS) | | `-cms, --checkpoint-min-step N` | minimum spacing between context checkpoints in tokens (default: 8192, 0 = no minimum)
(env: LLAMA_ARG_CHECKPOINT_MIN_SPACING_NT) | | `-cram, --cache-ram N` | set the maximum cache size in MiB (default: 8192, -1 - no limit, 0 - disable)[(more info)](https://github.com/ggml-org/llama.cpp/pull/16391)
(env: LLAMA_ARG_CACHE_RAM) | @@ -182,6 +185,9 @@ For the full list of features, please refer to [server's changelog](https://gith | `--image-min-tokens N` | minimum number of tokens each image can take, only used by vision models with dynamic resolution (default: read from model)
(env: LLAMA_ARG_IMAGE_MIN_TOKENS) | | `--image-max-tokens N` | maximum number of tokens each image can take, only used by vision models with dynamic resolution (default: read from model)
(env: LLAMA_ARG_IMAGE_MAX_TOKENS) | | `--mtmd-batch-max-tokens N` | maximum number of image tokens per batch when encoding images (default: 1024)
(env: LLAMA_ARG_MTMD_BATCH_MAX_TOKENS) | +| `--video-fps N` | target video frame rate (default: 4.0)
(env: LLAMA_ARG_VIDEO_FPS) | +| `--video-timestamp-interval N` | interval in milliseconds between text timestamps (default: 5000)
(env: LLAMA_ARG_VIDEO_TIMESTAMP_INTERVAL) | +| `--video-ffmpeg-dir DIR` | path to the directory containing ffmpeg and ffprobe (default: search in PATH)
(env: LLAMA_ARG_VIDEO_FFMPEG_DIR) | | `-a, --alias STRING` | set model name aliases, comma-separated (to be used by API)
(env: LLAMA_ARG_ALIAS) | | `--tags STRING` | set model tags, comma-separated (informational, not used for routing)
(env: LLAMA_ARG_TAGS) | | `--embd-normalize N` | normalisation for embeddings (default: 2) (-1=none, 0=max absolute int16, 1=taxicab, 2=euclidean, >2=p-norm) | @@ -256,6 +262,8 @@ For the full list of features, please refer to [server's changelog](https://gith | `--spec-draft-n-cpu-moe, --spec-draft-ncmoe, -ncmoed, --n-cpu-moe-draft N` | keep the Mixture of Experts (MoE) weights of the first N layers in the CPU for the draft model
(env: LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE) | | `--spec-draft-n-max N` | number of tokens to draft for speculative decoding (default: 3)
(env: LLAMA_ARG_SPEC_DRAFT_N_MAX) | | `--spec-draft-n-min N` | minimum number of draft tokens to use for speculative decoding (default: 0)
(env: LLAMA_ARG_SPEC_DRAFT_N_MIN) | +| `--spec-synth-len L` | target mean synthetic acceptance length, including the target token (benchmarking only)
(env: LLAMA_ARG_SPEC_SYNTH_LEN) | +| `--spec-synth-rates P0,P1,...` | comma-separated unconditional per-position synthetic acceptance probabilities (benchmarking only)
(env: LLAMA_ARG_SPEC_SYNTH_RATES) | | `--spec-draft-p-split, --draft-p-split P` | speculative decoding split probability (default: 0.10)
(env: LLAMA_ARG_SPEC_DRAFT_P_SPLIT) | | `--spec-draft-p-min, --draft-p-min P` | minimum speculative decoding probability (greedy) (default: 0.00)
(env: LLAMA_ARG_SPEC_DRAFT_P_MIN) | | `--spec-draft-backend-sampling, --no-spec-draft-backend-sampling` | offload draft sampling to the backend (default: enabled)
(env: LLAMA_ARG_SPEC_DRAFT_BACKEND_SAMPLING) | diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index 7997d4016a6..c30955e89f0 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -910,12 +910,17 @@ size_t validate_utf8(const std::string& text) { return len; } -server_tokens process_mtmd_prompt(mtmd_context * mctx, const std::string & prompt, const std::vector & files, bool is_placeholder) { +server_tokens process_mtmd_prompt( + mtmd_context * mctx, + const std::string & prompt, + const std::vector & files, + const mtmd_helper_init_opt & init_opt, + bool is_placeholder) { // these will be freed upon going out of scope mtmd::bitmaps bitmaps; std::vector videos; for (auto & file : files) { - auto out = mtmd_helper_bitmap_init_from_buf(mctx, file.data(), file.size(), is_placeholder); + auto out = mtmd_helper_bitmap_init_from_buf(mctx, file.data(), file.size(), is_placeholder, init_opt); if (!out.bitmap) { throw std::runtime_error("Failed to load image or audio file"); } @@ -956,7 +961,7 @@ server_tokens process_mtmd_prompt(mtmd_context * mctx, const std::string & promp * - "prompt": [12, 34, "string", 56, 78] * - "prompt": { "prompt_string": "string", "multimodal_data": [ "base64" ] } */ -static server_tokens tokenize_input_subprompt(const llama_vocab * vocab, mtmd_context * mctx, const json & json_prompt, bool add_special, bool parse_special) { +static server_tokens tokenize_input_subprompt(const llama_vocab * vocab, mtmd_context * mctx, const json & json_prompt, bool add_special, bool parse_special, const mtmd_helper_init_opt & init_opt) { constexpr char JSON_STRING_PROMPT_KEY[] = "prompt_string"; constexpr char JSON_MTMD_DATA_KEY[] = "multimodal_data"; const bool has_mtmd = mctx != nullptr; @@ -979,7 +984,7 @@ static server_tokens tokenize_input_subprompt(const llama_vocab * vocab, mtmd_co for (const auto & entry : json_prompt.at(JSON_MTMD_DATA_KEY)) { files.push_back(base64_decode(entry)); } - return process_mtmd_prompt(mctx, json_prompt.at(JSON_STRING_PROMPT_KEY), files); + return process_mtmd_prompt(mctx, json_prompt.at(JSON_STRING_PROMPT_KEY), files, init_opt); } else { // Not multimodal, but contains a subobject. llama_tokens tmp = tokenize_mixed(vocab, json_prompt.at(JSON_STRING_PROMPT_KEY), add_special, parse_special); @@ -990,15 +995,15 @@ static server_tokens tokenize_input_subprompt(const llama_vocab * vocab, mtmd_co } } -std::vector tokenize_input_prompts(const llama_vocab * vocab, mtmd_context * mctx, const json & json_prompt, bool add_special, bool parse_special) { +std::vector tokenize_input_prompts(const llama_vocab * vocab, mtmd_context * mctx, const json & json_prompt, bool add_special, bool parse_special, const mtmd_helper_init_opt & init_opt) { std::vector result; if (json_prompt.is_array() && !json_is_array_and_contains_numbers(json_prompt)) { result.reserve(json_prompt.size()); for (const auto & p : json_prompt) { - result.push_back(tokenize_input_subprompt(vocab, mctx, p,add_special, parse_special)); + result.push_back(tokenize_input_subprompt(vocab, mctx, p, add_special, parse_special, init_opt)); } } else { - result.push_back(tokenize_input_subprompt(vocab, mctx, json_prompt, add_special, parse_special)); + result.push_back(tokenize_input_subprompt(vocab, mctx, json_prompt, add_special, parse_special, init_opt)); } if (result.empty()) { throw std::runtime_error("\"prompt\" must not be empty"); @@ -1787,7 +1792,8 @@ server_tokens format_prompt_rerank( const struct llama_vocab * vocab, mtmd_context * mctx, const std::string & query, - const std::string & doc) { + const std::string & doc, + const mtmd_helper_init_opt & init_opt) { server_tokens result = {}; const char * rerank_prompt = llama_model_chat_template(model, "rerank"); @@ -1796,12 +1802,12 @@ server_tokens format_prompt_rerank( std::string prompt = rerank_prompt; string_replace_all(prompt, "{query}" , query); string_replace_all(prompt, "{document}", doc ); - server_tokens tokens = tokenize_input_subprompt(vocab, mctx, prompt, false, true); + server_tokens tokens = tokenize_input_subprompt(vocab, mctx, prompt, false, true, init_opt); result.push_back(tokens); } else { // Get EOS token - use SEP token as fallback if EOS is not available - server_tokens query_tokens = tokenize_input_subprompt(vocab, mctx, query, false, false); - server_tokens doc_tokens = tokenize_input_subprompt(vocab, mctx, doc, false, false); + server_tokens query_tokens = tokenize_input_subprompt(vocab, mctx, query, false, false, init_opt); + server_tokens doc_tokens = tokenize_input_subprompt(vocab, mctx, doc, false, false, init_opt); llama_token eos_token = llama_vocab_eos(vocab); if (eos_token == LLAMA_TOKEN_NULL) { eos_token = llama_vocab_sep(vocab); diff --git a/tools/server/server-common.h b/tools/server/server-common.h index f8ea82ef4cf..6c681a2cf56 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -5,6 +5,7 @@ #include "llama.h" #include "chat.h" #include "mtmd.h" +#include "mtmd-helper.h" #include "json.h" @@ -269,7 +270,12 @@ size_t validate_utf8(const std::string& text); // process mtmd prompt, return the server_tokens containing both text tokens and media chunks // if is_placeholder is true, the media chunk will be treated as placeholder for counting tokens; the output tokens are not usable for actual inference (e.g. for submitting a task to server_queue) -server_tokens process_mtmd_prompt(mtmd_context * mctx, const std::string & prompt, const std::vector & files, bool is_placeholder = false); +server_tokens process_mtmd_prompt( + mtmd_context * mctx, + const std::string & prompt, + const std::vector & files, + const mtmd_helper_init_opt & init_opt, + bool is_placeholder = false); /** * break the input "prompt" object into multiple prompt if needed, then tokenize them @@ -289,7 +295,8 @@ std::vector tokenize_input_prompts( mtmd_context * mctx, const json & json_prompt, bool add_special, - bool parse_special); + bool parse_special, + const mtmd_helper_init_opt & init_opt); // // OAI utils @@ -538,7 +545,8 @@ server_tokens format_prompt_rerank( const struct llama_vocab * vocab, mtmd_context * mctx, const std::string & query, - const std::string & doc); + const std::string & doc, + const mtmd_helper_init_opt & init_opt); // simple implementation of a pipe // used for streaming data between threads diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a9edbd7be8b..f5477356d61 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -23,6 +23,7 @@ #include #include #include +#include #include #include @@ -51,6 +52,50 @@ static common_speculative_output_limits server_output_limits(const common_params return result; } +// synthetic draft verification for benchmarking - accept draft tokens at random instead of by match with the target +// on replay the draft was already accepted before a context checkpoint restore, so repeat the same decisions +static std::vector server_sample_and_accept_synth( + common_sampler * smpl, + llama_context * ctx, + const std::vector & idxs, + const llama_tokens & draft, + const std::vector & synth_probs, + std::mt19937 & rng, + bool is_replay) { + GGML_ASSERT(idxs.size() == draft.size() + 1); + GGML_ASSERT(synth_probs.size() >= draft.size()); + + std::vector result; + result.reserve(idxs.size()); + + const llama_vocab * vocab = llama_model_get_vocab(llama_get_model(ctx)); + std::uniform_real_distribution dist(0.0, 1.0); + for (size_t i = 0; i < draft.size(); ++i) { + const llama_token id = common_sampler_sample(smpl, ctx, idxs[i]); + const bool accept = is_replay || dist(rng) < synth_probs[i]; + // do not accept a drafted EOG token - it would end the generation early + // on replay the last token is from the target and can be EOG, so skip this check + if (accept && (is_replay || !llama_vocab_is_eog(vocab, draft[i]))) { + // synthetic draft tokens do not advance grammar or reasoning state + // the last replay token is from the target and must advance both + const bool is_replay_target = is_replay && i + 1 == draft.size(); + common_sampler_accept(smpl, draft[i], is_replay_target); + result.push_back(draft[i]); + continue; + } + + common_sampler_accept(smpl, id, true); + result.push_back(id); + return result; + } + + const llama_token id = common_sampler_sample(smpl, ctx, idxs[draft.size()]); + common_sampler_accept(smpl, id, true); + result.push_back(id); + + return result; +} + // state diagram: https://github.com/ggml-org/llama.cpp/pull/9283 enum slot_state { SLOT_STATE_IDLE, @@ -211,6 +256,7 @@ struct server_slot { std::vector spec_i_batch; common_prompt_checkpoint spec_ckpt; bool spec_is_replay = false; + std::mt19937 spec_synth_rng; // TODO: move members that belong to the task (such as `generated_text`, `has_new_line`) to task_results_state // see https://github.com/ggml-org/llama.cpp/pull/18283#issuecomment-3710175837 @@ -794,6 +840,8 @@ struct server_context_impl { llama_model * model_tgt = nullptr; mtmd_context * mctx = nullptr; + // note: video_params.ffmpeg_bin_dir points into params_base, which outlives this struct + mtmd_helper_init_opt init_opt = mtmd_helper_init_opt_default(); const llama_vocab * vocab = nullptr; server_queue queue_tasks; @@ -1118,6 +1166,11 @@ struct server_context_impl { } SRV_INF("loaded multimodal model, '%s'\n", mmproj_path.c_str()); + init_opt.video_params.fps_target = params_base.video_fps; + init_opt.video_params.timestamp_interval_ms = params_base.video_timestamp_interval_ms; + init_opt.video_params.ffmpeg_bin_dir = params_base.video_ffmpeg_bin_dir.empty() + ? nullptr : params_base.video_ffmpeg_bin_dir.c_str(); + if (params_base.ctx_shift) { params_base.ctx_shift = false; SRV_WRN("%s\n", "ctx_shift is not supported by multimodal, it will be disabled"); @@ -1155,10 +1208,31 @@ struct server_context_impl { const int n_ctx_train = llama_model_n_ctx_train(model_tgt); - int n_ctx_slot = llama_n_ctx_seq(ctx_tgt); - if (n_ctx_slot > n_ctx_train) { - SRV_WRN("the slot context (%d) exceeds the training context of the model (%d) - capping\n", n_ctx_slot, n_ctx_train); - n_ctx_slot = n_ctx_train; + { + // note: the capping itself is done in n_ctx_slot(), here we only report it + const int n_ctx_seq = llama_n_ctx_seq(ctx_tgt); + + if (params_base.kv_unified_per_slot > 0) { + if (n_ctx_seq > params_base.kv_unified_per_slot) { + SRV_INF("capping per-slot context (%d) to --kv-unified-per-slot (%d)\n", + n_ctx_seq, params_base.kv_unified_per_slot); + } else if (params_base.kv_unified_per_slot > n_ctx_seq) { + // cap is above the per-slot pool capacity, so it can never bind + SRV_WRN( + "--kv-unified-per-slot (%d) exceeds the per-slot pool capacity (%d) - cap has no effect, " + "slots are limited to %d (raise the KV pool with -c, or unset -c to size it to " + "n_parallel * kv_unified_per_slot)\n", + params_base.kv_unified_per_slot, n_ctx_seq, n_ctx_seq); + } + } + + const int n_ctx_capped = params_base.kv_unified_per_slot > 0 ? + std::min(n_ctx_seq, params_base.kv_unified_per_slot) : n_ctx_seq; + + if (n_ctx_capped > n_ctx_train) { + SRV_WRN("the slot context (%d) exceeds the training context of the model (%d) - capping\n", + n_ctx_capped, n_ctx_train); + } } slots.clear(); @@ -1174,7 +1248,7 @@ struct server_context_impl { // setup slots SRV_INF("initializing, n_slots = %d, n_ctx_slot = %d, kv_unified = '%s'\n", - params_base.n_parallel, n_ctx_slot, params_base.kv_unified ? "true" : "false"); + params_base.n_parallel, n_ctx_slot(), params_base.kv_unified ? "true" : "false"); // initialize slots for (int i = 0; i < params_base.n_parallel; i++) { @@ -1187,6 +1261,9 @@ struct server_context_impl { spec.reset(common_speculative_init(params_base.speculative, params_base.n_parallel)); } catch (const std::exception & e) { SRV_ERR("failed to initialize speculative decoding context: %s\n", e.what()); + if (params_base.speculative.has_synth()) { + return false; + } } } @@ -1202,6 +1279,11 @@ struct server_context_impl { model_dft = nullptr; } + if (!spec && params_base.speculative.has_synth()) { + SRV_ERR("%s", "synthetic acceptance requires an initialized speculative decoding context\n"); + return false; + } + for (int i = 0; i < params_base.n_parallel; i++) { server_slot & slot = slots[i]; @@ -1210,7 +1292,7 @@ struct server_context_impl { slot.ctx_dft = ctx_dft; slot.mem.init(ctx_tgt, ctx_dft); slot.spec = spec.get(); - slot.n_ctx = n_ctx_slot; + slot.n_ctx = n_ctx_slot(); slot.mctx = mctx; slot.prompt.tokens.has_mtmd = mctx != nullptr; @@ -1710,6 +1792,13 @@ struct server_context_impl { SLT_TRC(slot, "sampler chain: %s\n", common_sampler_print(slot.smpl.get()).c_str()); SLT_TRC(slot, "sampler params: \n%s\n", task.params.sampling.print().c_str()); + + if (spec && !common_speculative_get_synth_probs(spec.get()).empty()) { + const uint32_t seed = task.params.sampling.seed == LLAMA_DEFAULT_SEED + ? std::random_device{}() + : task.params.sampling.seed; + slot.spec_synth_rng.seed(seed); + } } else { slot.smpl.reset(); } @@ -2134,9 +2223,9 @@ struct server_context_impl { try { auto & prompt = task.cli_prompt; if (mctx != nullptr) { - task.tokens = process_mtmd_prompt(mctx, prompt, task.cli_files); + task.tokens = process_mtmd_prompt(mctx, prompt, task.cli_files, init_opt); } else { - task.tokens = std::move(tokenize_input_prompts(vocab, mctx, prompt, true, true)[0]); + task.tokens = std::move(tokenize_input_prompts(vocab, mctx, prompt, true, true, init_opt)[0]); } task.cli_prompt.clear(); task.cli_files.clear(); @@ -3795,7 +3884,12 @@ struct server_context_impl { common_sampler_ptr smpl_save(common_sampler_clone(slot.smpl.get())); GGML_ASSERT(slot.spec_i_batch.size() == n_draft + 1); - auto accepted = common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft); + const auto & synth_probs = common_speculative_get_synth_probs(spec.get()); + auto accepted = synth_probs.empty() + ? common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft) + : server_sample_and_accept_synth( + slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft, + synth_probs, slot.spec_synth_rng, slot.spec_is_replay); slot.spec_i_batch.clear(); GGML_ASSERT(accepted.size() >= 1); @@ -3861,7 +3955,7 @@ struct server_context_impl { auto & n_accepted_per_pos = slot.n_accepted_per_pos; if (n_accepted_per_pos.empty()) { - n_accepted_per_pos.resize(common_speculative_n_max(¶ms_base.speculative), 0); + n_accepted_per_pos.resize(common_speculative_n_max(spec.get()), 0); } for (size_t i = 0; i < n_accepted && i < n_accepted_per_pos.size(); ++i) { n_accepted_per_pos[i]++; @@ -3902,8 +3996,15 @@ struct server_context_impl { }); } - int get_slot_n_ctx() { - return slots.back().n_ctx; + // context size of a single slot, capped by --kv-unified-per-slot and by the training context of the model + int n_ctx_slot() const { + int res = llama_n_ctx_seq(ctx_tgt); + + if (params_base.kv_unified_per_slot > 0) { + res = std::min(res, params_base.kv_unified_per_slot); + } + + return std::min(res, llama_model_n_ctx_train(model_tgt)); } server_response_reader get_response_reader() { @@ -4069,7 +4170,7 @@ server_context_meta server_context::get_meta() const { /* has_inp_audio */ impl->chat_params.allow_audio, /* has_inp_video */ impl->chat_params.allow_video, /* json_ui_settings */ impl->json_ui_settings, - /* slot_n_ctx */ impl->get_slot_n_ctx(), + /* slot_n_ctx */ impl->n_ctx_slot(), /* pooling_type */ llama_pooling_type(impl->ctx_tgt), /* chat_params */ impl->chat_params, @@ -4165,10 +4266,10 @@ std::unique_ptr server_routes::handle_completions_impl( if (res_type != TASK_RESPONSE_TYPE_NONE && ctx_server.mctx != nullptr) { // This is the case used by OAI compatible chat path with MTMD. TODO It can be moved to the path below. - inputs.push_back(process_mtmd_prompt(ctx_server.mctx, prompt.get(), files)); + inputs.push_back(process_mtmd_prompt(ctx_server.mctx, prompt.get(), files, ctx_server.init_opt)); } else { // Everything else, including multimodal completions. - inputs = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, true, true); + inputs = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, true, true, ctx_server.init_opt); } // tasks.reserve(inputs.size()); // TODO: this is inaccurate due to child tasks @@ -4752,7 +4853,7 @@ void server_routes::init_routes() { data["input_extra"] = input_extra; // default to empty array if it's not exist std::string prompt = json_value(data, "prompt", std::string()); - std::vector tokenized_prompts = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, false, true); + std::vector tokenized_prompts = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, false, true, ctx_server.init_opt); SRV_DBG("creating infill tasks, n_prompts = %d\n", (int) tokenized_prompts.size()); data["prompt"] = format_prompt_infill( ctx_server.vocab, @@ -4816,7 +4917,7 @@ void server_routes::init_routes() { }; this->post_chat_completions_tok = [this](const server_http_req & req) { - return handle_count_tokens(ctx_server.vocab, ctx_server.mctx, req, TASK_RESPONSE_TYPE_OAI_CHAT); + return handle_count_tokens(ctx_server.vocab, ctx_server.mctx, ctx_server.init_opt, req, TASK_RESPONSE_TYPE_OAI_CHAT); }; this->post_control = [this](const server_http_req & req) { @@ -4875,7 +4976,7 @@ void server_routes::init_routes() { }; this->post_responses_tok_oai = [this](const server_http_req & req) { - return handle_count_tokens(ctx_server.vocab, ctx_server.mctx, req, TASK_RESPONSE_TYPE_OAI_RESP); + return handle_count_tokens(ctx_server.vocab, ctx_server.mctx, ctx_server.init_opt, req, TASK_RESPONSE_TYPE_OAI_RESP); }; this->post_transcriptions_oai = [this](const server_http_req & req) { @@ -4925,7 +5026,7 @@ void server_routes::init_routes() { }; this->post_anthropic_count_tokens = [this](const server_http_req & req) { - return handle_count_tokens(ctx_server.vocab, ctx_server.mctx, req, TASK_RESPONSE_TYPE_ANTHROPIC); + return handle_count_tokens(ctx_server.vocab, ctx_server.mctx, ctx_server.init_opt, req, TASK_RESPONSE_TYPE_ANTHROPIC); }; // same with handle_chat_completions, but without inference part @@ -5058,7 +5159,7 @@ void server_routes::init_routes() { std::vector tasks; tasks.reserve(documents.size()); for (size_t i = 0; i < documents.size(); i++) { - auto tmp = format_prompt_rerank(ctx_server.model_tgt, ctx_server.vocab, ctx_server.mctx, query, documents[i]); + auto tmp = format_prompt_rerank(ctx_server.model_tgt, ctx_server.vocab, ctx_server.mctx, query, documents[i], ctx_server.init_opt); server_task task = server_task(SERVER_TASK_TYPE_RERANK); task.id = rd.get_new_id(); task.tokens = std::move(tmp); @@ -5296,7 +5397,7 @@ std::unique_ptr server_routes::handle_embeddings_impl(cons } } - auto tokenized_prompts = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, true, true); + auto tokenized_prompts = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, true, true, ctx_server.init_opt); for (const auto & tokens : tokenized_prompts) { // this check is necessary for models that do not add BOS token to the input if (tokens.empty()) { @@ -5357,7 +5458,7 @@ std::unique_ptr server_routes::handle_embeddings_impl(cons return res; } -std::unique_ptr server_routes::handle_count_tokens(const llama_vocab * vocab, mtmd_context * mctx, const server_http_req & req, task_response_type res_type) { +std::unique_ptr server_routes::handle_count_tokens(const llama_vocab * vocab, mtmd_context * mctx, const mtmd_helper_init_opt & init_opt, const server_http_req & req, task_response_type res_type) { auto res = create_response(); std::vector files; json body = json::parse(req.body); @@ -5395,7 +5496,7 @@ std::unique_ptr server_routes::handle_count_tokens(const l if (!prompt.is_string()) { throw std::runtime_error("for mtmd, input prompt must be a string."); } - n_tokens = process_mtmd_prompt(mctx, prompt.get(), files, true).size(); + n_tokens = process_mtmd_prompt(mctx, prompt.get(), files, init_opt, true).size(); } else { n_tokens = tokenize_mixed(vocab, prompt, true, true).size(); } diff --git a/tools/server/server-context.h b/tools/server/server-context.h index 5d464b8e8cb..0acbbffa9e1 100644 --- a/tools/server/server-context.h +++ b/tools/server/server-context.h @@ -169,7 +169,7 @@ struct server_routes { std::unique_ptr handle_slots_restore(const server_http_req & req, int id_slot); std::unique_ptr handle_slots_erase(const server_http_req &, int id_slot); std::unique_ptr handle_embeddings_impl(const server_http_req & req, task_response_type res_type); - std::unique_ptr handle_count_tokens(const llama_vocab * vocab, mtmd_context * mctx, const server_http_req & req, task_response_type res_type); + std::unique_ptr handle_count_tokens(const llama_vocab * vocab, mtmd_context * mctx, const mtmd_helper_init_opt & init_opt, const server_http_req & req, task_response_type res_type); // using unique_ptr to allow late initialization of const std::unique_ptr meta; diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 5fe2729ba1b..22378b38c5e 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -157,6 +157,18 @@ int llama_server(common_params & params, int argc, char ** argv) { } } + // size the KV pool from --kv-unified-per-slot, unless the user pinned it with -c + // or with -c 0 for max context + const bool ctx_pool_auto_sized = params.kv_unified_per_slot > 0 && + params.n_ctx == 0 && + (uint32_t) params.fit_params_min_ctx != UINT32_MAX; + + if (ctx_pool_auto_sized) { + params.n_ctx = params.n_parallel * params.kv_unified_per_slot; + SRV_INF("--kv-unified-per-slot: sizing KV pool to n_parallel * kv_unified_per_slot = %d * %d = %d\n", params.n_parallel, + params.kv_unified_per_slot, params.n_ctx); + } + // for consistency between server router mode and single-model mode, we set the same model name as alias auto model_name = params.model.get_name(); if (params.model_alias.empty() && !model_name.empty()) { diff --git a/tools/server/tests/unit/test_speculative.py b/tools/server/tests/unit/test_speculative.py index 5837195006b..22b523954ec 100644 --- a/tools/server/tests/unit/test_speculative.py +++ b/tools/server/tests/unit/test_speculative.py @@ -52,6 +52,18 @@ def test_with_and_without_draft(): assert tokens_no_draft == tokens_draft + server.stop() + create_server() + assert server.spec_draft_n_max is not None + server.spec_synth_rates = [0.0] * server.spec_draft_n_max + server.start() + res = server.make_request("POST", "/completion", data=request) + + assert res.status_code == 200 + assert res.body["timings"]["draft_n"] > 0 + assert res.body["timings"]["draft_n_accepted"] == 0 + assert res.body["tokens"] == tokens_no_draft + def test_different_draft_min_draft_max(): global server @@ -80,6 +92,66 @@ def test_different_draft_min_draft_max(): last_content = res.body["content"] +def test_synth_is_deterministic(): + global server + assert server.spec_draft_n_max is not None + server.spec_synth_rates = [0.75 ** (i + 1) for i in range(server.spec_draft_n_max)] + server.start() + + request = { + "prompt": "I believe the meaning of life is", + "temperature": 0.2, + "top_k": 5, + "seed": 4242, + "n_predict": 32, + } + responses = [server.make_request("POST", "/completion", data=request) for _ in range(2)] + + for res in responses: + assert res.status_code == 200 + assert res.body["timings"]["draft_n"] > 0 + assert responses[0].body["timings"]["draft_n"] == responses[1].body["timings"]["draft_n"] + assert responses[0].body["timings"]["draft_n_accepted"] == responses[1].body["timings"]["draft_n_accepted"] + + +def test_synth_ignores_target_tokens(): + global server + assert server.spec_draft_n_max is not None + server.spec_synth_rates = [1.0] * server.spec_draft_n_max + server.start() + + res = server.make_request("POST", "/completion", data={ + "prompt": "I believe the meaning of life is", + "temperature": 0.0, + "seed": 4242, + "n_predict": 32, + }) + + assert res.status_code == 200 + assert res.body["timings"]["draft_n"] > 0 + assert res.body["timings"]["draft_n_accepted"] == res.body["timings"]["draft_n"] + + res = server.make_request("POST", "/completion", data={ + "prompt": "I believe the meaning of life is", + "temperature": 0.0, + "seed": 4242, + "n_predict": 6, + "grammar": 'root ::= "a"{5,5}', + }) + assert res.status_code == 200, res.body + + res = server.make_request("POST", "/completion", data={ + "prompt": "Respond with only: OK", + "temperature": 0.0, + "seed": 4242, + "n_predict": 64, + "ignore_eos": True, + }) + assert res.status_code == 200, res.body + assert res.body["tokens_predicted"] == 64 + assert res.body["stop_type"] == "limit" + + def test_slot_ctx_not_exceeded(): global server server.n_ctx = 256 diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py index a0d2dfa3c59..5a4f31a53e0 100644 --- a/tools/server/tests/utils.py +++ b/tools/server/tests/utils.py @@ -99,6 +99,8 @@ class ServerProcess: spec_type: str | None = None spec_draft_n_min: int | None = None spec_draft_n_max: int | None = None + spec_synth_len: float | None = None + spec_synth_rates: List[float] | None = None no_ui: bool | None = None jinja: bool | None = None reasoning_format: Literal['deepseek', 'none', 'nothink'] | None = None @@ -245,6 +247,11 @@ def start(self, timeout_seconds: int = DEFAULT_HTTP_TIMEOUT) -> None: server_args.extend(["--spec-draft-n-max", self.spec_draft_n_max]) if self.spec_draft_n_min: server_args.extend(["--spec-draft-n-min", self.spec_draft_n_min]) + if self.spec_synth_len is not None: + server_args.extend(["--spec-synth-len", self.spec_synth_len]) + if self.spec_synth_rates is not None: + rates = ",".join(str(rate) for rate in self.spec_synth_rates) + server_args.extend(["--spec-synth-rates", rates]) if self.no_ui: server_args.append("--no-ui") if self.no_models_autoload: diff --git a/tools/tts/tts.cpp b/tools/tts/tts.cpp index 368123baf53..6fd1936324c 100644 --- a/tools/tts/tts.cpp +++ b/tools/tts/tts.cpp @@ -103,7 +103,7 @@ int main(int argc, char ** argv) { mtmd::bitmap_ptr speaker_bitmap; if (!params.tts_speaker_file.empty()) { - auto wrapper = mtmd_helper_bitmap_init_from_file(mctx.get(), params.tts_speaker_file.c_str(), false); + auto wrapper = mtmd_helper_bitmap_init_from_file(mctx.get(), params.tts_speaker_file.c_str(), false, mtmd_helper_init_opt_default()); if (!wrapper.bitmap) { LOG_ERR("failed to load speaker file %s\n", params.tts_speaker_file.c_str()); return 1; diff --git a/tools/ui/src/lib/components/app/badges/BadgesModality.svelte b/tools/ui/src/lib/components/app/badges/BadgesModality.svelte index 4eb3e7838d1..83b1b46affc 100644 --- a/tools/ui/src/lib/components/app/badges/BadgesModality.svelte +++ b/tools/ui/src/lib/components/app/badges/BadgesModality.svelte @@ -1,5 +1,5 @@ -{#each modalities as modality (modality)} - {#if modality === ModelModality.VISION || modality === ModelModality.AUDIO || modality === ModelModality.VIDEO} - - {#if modality === ModelModality.VISION} - + const shownModalities = [ModelModality.VISION, ModelModality.AUDIO, ModelModality.VIDEO] as const; - Vision (Image) - {:else if modality === ModelModality.VIDEO} - {/each} diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatForm.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatForm.svelte index 9152935f073..893f8077dd6 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatForm.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatForm.svelte @@ -8,7 +8,8 @@ ChatFormInputFileInputInvisible, ChatFormMcpResourcesList, ChatFormPickers, - DialogMcpResourcesBrowser + DialogMcpResourcesBrowser, + DialogMcpServers } from '$lib/components/app'; import { CLIPBOARD_CONTENT_QUOTE_PREFIX, @@ -22,7 +23,8 @@ FileExtensionText, KeyboardKey, MimeTypeText, - SpecialFileType + SpecialFileType, + ToolSource } from '$lib/enums'; import { useChatFormPickers } from '$lib/hooks/use-chat-form-pickers.svelte'; import { @@ -72,7 +74,6 @@ disabled?: boolean; isLoading?: boolean; placeholder?: string; - showMcpPromptButton?: boolean; showAddButton?: boolean; showModelSelector?: boolean; @@ -102,7 +103,6 @@ onValueChange, placeholder = 'Type a message...', showAddButton = true, - showMcpPromptButton = false, showModelSelector = true, uploadedFiles = $bindable([]), value = $bindable('') @@ -151,9 +151,18 @@ getServerHome: () => toolsStore.serverHome ?? null, getShowModelSelector: () => showModelSelector, getValue: () => value, - hasCwdTools: () => toolsStore.hasEnabledCwdTools, - hasPrompts: () => - mcpStore.hasPromptsCapability(conversationsStore.preferences.getAllMcpServerOverrides()), + hasCwdTools: () => conversationsStore.preferences.hasEnabledCwdTools(), + // policy-aware, same rule as the agentic flow: MCP category on and at + // least one globally-enabled server whose group key is not disabled + hasPrompts: () => { + const prefs = conversationsStore.preferences; + + if (!prefs.isCategoryEnabled(ToolSource.MCP)) return false; + + return mcpStore + .getServers() + .some((s) => s.enabled && prefs.isServerToolsEnabled(s.id) && s.url.trim()); + }, openModelSelector: () => chatFormActionsRef?.openModelSelector(), setCaretOffset: (offset) => inputRef?.setCaretOffset(offset), setValue: (v) => { @@ -183,6 +192,9 @@ let isResourceDialogOpen = $state(false); let preSelectedResourceUri = $state(undefined); + // MCP Servers Dialog State + let isMcpServersDialogOpen = $state(false); + let currentConfig = $derived(settingsStore.config); let pasteLongTextToFileLength = $derived.by(() => { @@ -616,8 +628,7 @@ isReasoning={chatStore.isReasoning} {isRecording} onFileUpload={handleFileUpload} - onMcpPromptClick={showMcpPromptButton ? () => pickers.openPromptPicker() : undefined} - onMcpResourcesClick={() => (isResourceDialogOpen = true)} + onMcpSettingsClick={() => (isMcpServersDialogOpen = true)} onMicClick={handleMicClick} {onStop} onSystemPromptClick={() => onSystemPromptClick?.({ files: uploadedFiles, message: value })} @@ -630,7 +641,7 @@ - {#if toolsStore.hasEnabledCwdTools} + {#if conversationsStore.preferences.hasEnabledCwdTools()} + + diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddDropdown.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddDropdown.svelte index 02bfadb7e41..9c3a9e89122 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddDropdown.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddDropdown.svelte @@ -1,10 +1,6 @@
@@ -93,54 +92,32 @@ } }} > - - - - - - - + attachmentMenu.callbacks[AttachmentAction.FILE_UPLOAD]()} + > + + Add files - - - - {#each ATTACHMENT_FILE_ITEMS as item (item.id)} - {@const enabled = attachmentMenu.isItemEnabled(item.enabledWhen)} - {#if enabled} - attachmentMenu.callbacks[item.action]()} - > - - - {item.label} - - {:else if item.disabledTooltip} - - - {#snippet child({ props })} -
- - - - {item.label} - -
- {/snippet} - - - -

{item.disabledTooltip}

- - - {/if} - {/each} - - + + {#if supportedModalities.length > 0} + + {#each supportedModalities as modality (modality.label)} + + + + + + +

{modality.label}

+ + + {/each} + + {/if} + + - - - {#if chatFormActions.hasMcpPromptsSupport} - - - - - - MCP Prompt - - {/if} - - {#if chatFormActions.hasMcpResourcesSupport} - - + + - MCP Resources - - {/if} + MCP Servers +
diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddMcpServersSubmenu.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddMcpServersSubmenu.svelte deleted file mode 100644 index bceb43d2abd..00000000000 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddMcpServersSubmenu.svelte +++ /dev/null @@ -1,152 +0,0 @@ - - - - - - - - MCP Servers - - - - {#if hasMcpServers} - -
- {#each filteredMcpServers as server (server.id)} - {@const healthState = mcpStore.getHealthCheckState(server.id)} - {@const hasError = healthState.status === HealthCheckStatus.ERROR} - {@const isEnabledForChat = isServerEnabledForChat(server.id)} - {@const displayName = getServerLabel(server)} - {@const faviconUrl = mcpStore.getServerFavicon(server.id)} - - - {/each} -
- - {#snippet footer()} - - - - Manage MCP Servers - - {/snippet} - - {:else} -
- No MCP servers configured -
- - - - - - - Add MCP Servers - - {/if} - - - diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddReasoningSubmenu.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddReasoningSubmenu.svelte index 1b6fc4b0209..197d6c2e590 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddReasoningSubmenu.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddReasoningSubmenu.svelte @@ -8,70 +8,68 @@ const reasoning = useReasoningMenu(); -{#if reasoning.modelSupportsThinking} - - - {#if reasoning.thinkingEnabled} - - {:else if reasoning.isOff} - - {:else} - - {/if} + + + {#if reasoning.isReasoningActive} + + {:else if reasoning.isOff} + + {:else} + + {/if} - - Reasoning + + Reasoning - - {reasoning.currentEffort} - + + {reasoning.currentEffort} - + + - - {#each reasoning.levels as level (level.value)} - {@const tokenLabel = reasoning.tokenLabel(level)} - reasoning.select(level)} - > - {#if reasoning.isSelected(level)} - - {:else} -
- {/if} + + {#each reasoning.levels as level (level.value)} + {@const tokenLabel = reasoning.tokenLabel(level)} + reasoning.select(level)} + > + {#if reasoning.isSelected(level)} + + {:else} +
+ {/if} - {level.label} + {level.label} - {#if tokenLabel} - - {tokenLabel} - - {/if} + {#if tokenLabel} + + {tokenLabel} + + {/if} - {#if level.hasInfo} - - - - + {#if level.hasInfo} + + + + - -

Maximum reasoning effort with extended context usage

- - - {/if} - - {/each} - - -{/if} + +

Maximum reasoning effort with extended context usage

+ + + {/if} + + {/each} + + diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddSheet.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddSheet.svelte index 2f69dc96de7..acd0f4d2116 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddSheet.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddSheet.svelte @@ -1,18 +1,18 @@
@@ -194,80 +186,15 @@ - (mcpExpanded = open)} open={mcpExpanded}> - - {#if mcpExpanded} - - {:else} - - {/if} - - - - MCP Servers - - - {mcpServers.length} server{mcpServers.length !== 1 ? 's' : ''} - - - - -
- {#each mcpServers as server (server.id)} - {@const healthState = mcpStore.getHealthCheckState(server.id)} - {@const hasError = healthState.status === HealthCheckStatus.ERROR} - {@const displayName = mcpStore.getServerLabel(server)} - {@const faviconUrl = mcpStore.getServerFavicon(server.id)} - {@const isEnabled = conversationsStore.preferences.isMcpServerEnabledForChat( - server.id - )} - - - {/each} +
- - + System Message + {#if toolsPanel.totalToolCount > 0} (toolsExpanded = open)} open={toolsExpanded}> @@ -289,40 +216,12 @@
- {#each toolsPanel.activeGroups as group (group.key)} - {@const checked = toolsPanel.isGroupChecked(group)} - {@const enabledCount = toolsPanel.getEnabledToolCount(group)} - {@const favicon = toolsPanel.getFavicon(group)} - - + {#each toolsPanel.mcpGroups as group (group.key)} + {@render sheetGroupRow(group)} {/each}
@@ -331,38 +230,55 @@ - - {#if chatFormActions.hasMcpPromptsSupport} - - {/if} - - {#if chatFormActions.hasMcpResourcesSupport} - - {/if}
+ +{#snippet sheetGroupRow(group: ToolGroup)} + {@const checkState = toolsPanel.getGroupCheckState(group)} + {@const enabledCount = toolsPanel.getEnabledToolCount(group)} + {@const favicon = toolsPanel.getFavicon(group)} + {@const groupDisabled = toolsPanel.isGroupDisabled(group)} + + +{/snippet} diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte index 40fed27c70a..f4954417148 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte @@ -7,6 +7,7 @@ import { CLI_FLAGS, ICON_CLASS_DEFAULT } from '$lib/constants'; import { useToolsPanel } from '$lib/hooks/use-tools-panel.svelte'; import { mcpStore, toolsStore } from '$lib/stores'; + import type { ToolGroup } from '$lib/types'; const toolsPanel = useToolsPanel(); const hasMcpServersAvailable = $derived(mcpStore.getServers().length > 0); @@ -62,95 +63,108 @@ {/if} {:else}
- {#each toolsPanel.activeGroups as group (group.key)} - {@const isExpanded = toolsPanel.expandedGroups.has(group.key)} - {@const checked = toolsPanel.isGroupChecked(group)} - {@const favicon = toolsPanel.getFavicon(group)} - - toolsPanel.toggleGroupExpanded(group.key)} - open={isExpanded} - > -
- - {#if isExpanded} - - {:else} - - {/if} - - - {#if favicon} - { - (e.currentTarget as HTMLImageElement).style.display = 'none'; - }} - src={favicon} - /> - {/if} - - {group.label} - - - - {toolsPanel.getEnabledToolCount(group)}/{group.tools.length} - - - - - - {#snippet child({ props })} - toolsPanel.toggleGroupByKey(group.key)} - /> - {/snippet} - - - -

- {checked ? 'Disable' : 'Enable'} - {group.tools.length} tool{group.tools.length !== 1 ? 's' : ''} -

- - -
- - -
- {#each group.tools as entry (entry.key)} - {@const enabled = toolsStore.isToolEnabled(entry.key)} - - {/each} -
- - + {#each toolsPanel.categoryGroups as group (group.key)} + {@render groupRow(group)} + {/each} + + {#each toolsPanel.mcpGroups as group (group.key)} + {@render groupRow(group)} {/each}
{/if} + +{#snippet groupRow(group: ToolGroup)} + {@const isExpanded = toolsPanel.expandedGroups.has(group.key)} + {@const checkState = toolsPanel.getGroupCheckState(group)} + {@const favicon = toolsPanel.getFavicon(group)} + {@const groupDisabled = toolsPanel.isGroupDisabled(group)} + + toolsPanel.toggleGroupExpanded(group.key)} + open={isExpanded} + > +
+ + {#if isExpanded} + + {:else} + + {/if} + + + {#if favicon} + { + (e.currentTarget as HTMLImageElement).style.display = 'none'; + }} + src={favicon} + /> + {/if} + + {group.label} + + + + {toolsPanel.getEnabledToolCount(group)}/{group.tools.length} + + + + + + {#snippet child({ props })} + toolsPanel.toggleGroupByKey(group.key)} + /> + {/snippet} + + + +

+ {checkState.checked ? 'Disable' : 'Enable'} + {group.tools.length} tool{group.tools.length !== 1 ? 's' : ''} +

+ + +
+ + +
+ {#each group.tools as entry (entry.key)} + {@const enabled = toolsPanel.isToolEnabled(entry)} + {@const parentDisabled = toolsPanel.isToolParentDisabled(entry)} + + {/each} +
+ + +{/snippet} diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte index 97351b3a6cb..395f2cfbe17 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte @@ -1,6 +1,5 @@ - + diff --git a/tools/ui/src/lib/components/app/dialogs/DialogMcpServerAddNew.svelte b/tools/ui/src/lib/components/app/dialogs/DialogMcpServerAddNew.svelte index ae8b24cb2d3..bc28a754c61 100644 --- a/tools/ui/src/lib/components/app/dialogs/DialogMcpServerAddNew.svelte +++ b/tools/ui/src/lib/components/app/dialogs/DialogMcpServerAddNew.svelte @@ -10,7 +10,7 @@ RECOMMENDED_MCP_SERVERS } from '$lib/constants'; import { BooleanString, HealthCheckStatus } from '$lib/enums'; - import { conversationsStore, mcpStore } from '$lib/stores'; + import { mcpStore } from '$lib/stores'; import { canonicalizeServerUrl, parseHeadersToArray, uuid } from '$lib/utils'; interface Props { @@ -234,8 +234,6 @@ useProxy: newServerUseProxy }); - conversationsStore.preferences.setMcpServerOverride(newServerId, true); - handleOpenChange(false); } @@ -246,7 +244,7 @@ - + Add New MCP Server diff --git a/tools/ui/src/lib/components/app/dialogs/DialogMcpServers.svelte b/tools/ui/src/lib/components/app/dialogs/DialogMcpServers.svelte new file mode 100644 index 00000000000..2eaaa0a6577 --- /dev/null +++ b/tools/ui/src/lib/components/app/dialogs/DialogMcpServers.svelte @@ -0,0 +1,33 @@ + + + + + + + + + MCP Servers + + + + + + diff --git a/tools/ui/src/lib/components/app/dialogs/DialogMermaidPreview.svelte b/tools/ui/src/lib/components/app/dialogs/DialogMermaidPreview.svelte index 09e53442ac6..e741373497a 100644 --- a/tools/ui/src/lib/components/app/dialogs/DialogMermaidPreview.svelte +++ b/tools/ui/src/lib/components/app/dialogs/DialogMermaidPreview.svelte @@ -14,6 +14,7 @@ diff --git a/tools/ui/src/lib/components/app/dialogs/DialogModelInformation.svelte b/tools/ui/src/lib/components/app/dialogs/DialogModelInformation.svelte index 811c24d6b79..e200c004e58 100644 --- a/tools/ui/src/lib/components/app/dialogs/DialogModelInformation.svelte +++ b/tools/ui/src/lib/components/app/dialogs/DialogModelInformation.svelte @@ -76,22 +76,19 @@ - - - - - Model Information - - Current model details and capabilities - - -
+ + + + +
+
+ Model Information + + Current model details and capabilities +
+ {#if isLoadingModels || isLoadingRouterProps}
Loading model information...
@@ -100,17 +97,15 @@ {@const modelMeta = firstModel.meta} {#if serverProps} - + +