diff --git a/.github/actions/windows-setup-rocm/action.yml b/.github/actions/windows-setup-rocm/action.yml index fd9f8e5a4168..aecbcf14f522 100644 --- a/.github/actions/windows-setup-rocm/action.yml +++ b/.github/actions/windows-setup-rocm/action.yml @@ -8,8 +8,26 @@ inputs: runs: using: "composite" steps: - - name: Setup ROCm - uses: ./.github/actions/install-exe - with: - url: https://download.amd.com/developer/eula/rocm-hub/AMD-Software-PRO-Edition-${{ inputs.version }}-Win11-For-HIP.exe - args: -install + - name: Install ROCm with Wheels + shell: pwsh + run: | + $ErrorActionPreference = "Stop" + write-host "Setting up Python virtual environment" + + # Create the venv directly at the cache location to avoid relocation issues + New-Item -Path "C:\TheRock\build" -ItemType Directory -Force | Out-Null + python -m venv C:\TheRock\build\.venv + & C:\TheRock\build\.venv\Scripts\Activate.ps1 + + write-host "Upgrading pip" + python -m pip install --upgrade pip + + write-host "Installing ROCm wheels for multi-arch support" + # Install ROCm wheels for multi-arch support (this may take several minutes) + python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{ inputs.version }}" + + # Pre-expand the devel tree so it is included in the cache + write-host "Initializing ROCm devel tree" + rocm-sdk init + if ($LASTEXITCODE -ne 0) { throw "rocm-sdk init failed with exit code $LASTEXITCODE" } + write-host "Completed ROCm wheel installation to C:\TheRock\build" diff --git a/.github/workflows/build-cache.yml b/.github/workflows/build-cache.yml index 327f71978bf1..e15fc5e0830f 100644 --- a/.github/workflows/build-cache.yml +++ b/.github/workflows/build-cache.yml @@ -123,8 +123,8 @@ jobs: runs-on: windows-2022 env: - # Make sure this is in sync with build.yml - HIPSDK_INSTALLER_VERSION: "26.Q1" + # Make sure this is in sync with release.yml and build-cuda-windows.yml + ROCM_VERSION: "7.14.0" steps: - name: Clone @@ -135,11 +135,11 @@ jobs: uses: actions/cache@v5 id: cache-rocm with: - path: C:\Program Files\AMD\ROCm - key: cache-gha-rocm-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ runner.os }} + path: C:\TheRock\build + key: rocm-wheels-${{ env.ROCM_VERSION }}-multi-arch-${{ runner.os }} - name: Setup ROCm if: steps.cache-rocm.outputs.cache-hit != 'true' uses: ./.github/actions/windows-setup-rocm with: - version: ${{ env.HIPSDK_INSTALLER_VERSION }} + version: ${{ env.ROCM_VERSION }} diff --git a/.github/workflows/build-cuda-windows.yml b/.github/workflows/build-cuda-windows.yml index 367a3a8546cf..ff900802f4a2 100644 --- a/.github/workflows/build-cuda-windows.yml +++ b/.github/workflows/build-cuda-windows.yml @@ -83,7 +83,7 @@ jobs: env: # Make sure this is in sync with build-cache.yml - HIPSDK_INSTALLER_VERSION: "26.Q1" + ROCM_VERSION: "7.14.0" strategy: matrix: @@ -97,36 +97,53 @@ jobs: id: checkout uses: actions/checkout@v6 - - name: Grab rocWMMA package - id: grab_rocwmma - run: | - curl -o rocwmma.deb "https://repo.radeon.com/rocm/apt/7.2.1/pool/main/r/rocwmma-dev/rocwmma-dev_2.2.0.70201-81~24.04_amd64.deb" - 7z x rocwmma.deb - 7z x data.tar - - - name: Use ROCm Installation Cache + - name: Cache ROCm Installation uses: actions/cache@v5 id: cache-rocm with: - path: C:\Program Files\AMD\ROCm - key: cache-gha-rocm-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ runner.os }} + path: C:\TheRock\build + key: rocm-wheels-${{ env.ROCM_VERSION }}-multi-arch-${{ runner.os }} - name: Setup ROCm if: steps.cache-rocm.outputs.cache-hit != 'true' uses: ./.github/actions/windows-setup-rocm with: - version: ${{ env.HIPSDK_INSTALLER_VERSION }} + version: ${{ env.ROCM_VERSION }} + + - name: Setup ROCm Environment + run: | + $ErrorActionPreference = "Stop" + + # Activate venv from cache or fresh install + & C:\TheRock\build\.venv\Scripts\Activate.ps1 + + # Expand the devel tree (idempotent; no-op if already done during install) + rocm-sdk init + if ($LASTEXITCODE -ne 0) { throw "rocm-sdk init failed with exit code $LASTEXITCODE" } + + # Get ROCm installation paths using the rocm-sdk CLI tool + $rocmPath = (rocm-sdk path --root) + if (-not $rocmPath) { throw "rocm-sdk path --root returned empty - devel package may not be installed" } + $rocmPath = $rocmPath.Trim() + $cmakePath = (rocm-sdk path --cmake).Trim() + $binPath = (rocm-sdk path --bin).Trim() + write-host "ROCm root: $rocmPath" + + echo "HIP_PATH=$rocmPath" >> $env:GITHUB_ENV + echo "CMAKE_PREFIX_PATH=$cmakePath" >> $env:GITHUB_ENV + echo "HIP_DEVICE_LIB_PATH=$rocmPath\lib\llvm\amdgcn\bitcode" >> $env:GITHUB_ENV + echo "HIP_PLATFORM=amd" >> $env:GITHUB_ENV + echo "LLVM_PATH=$rocmPath\lib\llvm" >> $env:GITHUB_ENV + echo "$binPath" >> $env:GITHUB_PATH + + # Keep venv in PATH for subsequent steps + echo "C:\TheRock\build\.venv\Scripts" >> $env:GITHUB_PATH - name: Verify ROCm id: verify run: | - # Find and test ROCm installation - $clangPath = Get-ChildItem 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | Select-Object -First 1 - if (-not $clangPath) { - Write-Error "ROCm installation not found" - exit 1 - } - & $clangPath.FullName --version + # Test the ROCm clang shipped in the installed wheel + & "${env:HIP_PATH}\lib\llvm\bin\clang.exe" --version - name: ccache uses: ggml-org/ccache-action@v1.2.21 @@ -134,28 +151,27 @@ jobs: # TODO: this build does not match the build in release.yml, so we use a different cache key # ideally, the builds should match, similar to the CUDA build above so that we would be able # to populate the ccache for the release with manual runs of this workflow - #key: release-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }} - key: cuda-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }} + #key: release-windows-2022-x64-hip-${{ env.ROCM_VERSION }}-${{ matrix.name }} + key: cuda-windows-2022-x64-hip-${{ env.ROCM_VERSION }}-${{ matrix.name }} - name: Build id: cmake_build run: | - $env:HIP_PATH=$(Resolve-Path 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | split-path | split-path) - $env:CMAKE_PREFIX_PATH="${env:HIP_PATH}" cmake -G "Unix Makefiles" -B build -S . ` - -DCMAKE_C_COMPILER="${env:HIP_PATH}\bin\clang.exe" ` - -DCMAKE_CXX_COMPILER="${env:HIP_PATH}\bin\clang++.exe" ` - -DCMAKE_CXX_FLAGS="-I$($PWD.Path.Replace('\', '/'))/opt/rocm-7.2.1/include/" ` + -DCMAKE_PREFIX_PATH="${env:HIP_PATH}" ` + -DCMAKE_C_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" ` + -DCMAKE_CXX_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang++.exe" ` + -DCMAKE_HIP_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" ` -DCMAKE_BUILD_TYPE=Release ` -DLLAMA_BUILD_BORINGSSL=ON ` - -DROCM_DIR="${env:HIP_PATH}" ` + -DHIP_PATH="${env:HIP_PATH}" ` -DGGML_HIP=ON ` - -DGPU_TARGETS="gfx1100" ` + -DGPU_TARGETS="gfx1100" ` -DGGML_RPC=ON cmake --build build -j ${env:NUMBER_OF_PROCESSORS} - name: ccache-clear uses: ./.github/actions/ccache-clear with: - #key: release-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }} - key: cuda-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }} + #key: release-windows-2022-x64-hip-${{ env.ROCM_VERSION }}-${{ matrix.name }} + key: cuda-windows-2022-x64-hip-${{ env.ROCM_VERSION }}-${{ matrix.name }} diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index c668930b0caa..7f47c72135f8 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -748,6 +748,133 @@ jobs: path: llama-bin-win-cpu-${{ matrix.arch }}.zip name: llama-bin-win-cpu-${{ matrix.arch }}.zip + windows-rocm: + needs: [check-release] + runs-on: windows-2022 + + strategy: + matrix: + include: + - ROCM_VERSION: "7.14.0" + gpu_targets: "gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201" + build: x64 + + steps: + - name: Clone + id: checkout + uses: actions/checkout@v6 + with: + fetch-depth: 0 + + - name: ccache + uses: ggml-org/ccache-action@v1.2.21 + with: + key: windows-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }} + evict-old-files: 1d + + - name: Cache ROCm Installation + id: cache-rocm + uses: actions/cache@v5 + with: + path: C:\TheRock\build + key: rocm-wheels-${{ matrix.ROCM_VERSION }}-multi-arch-${{ runner.os }} + + - name: Setup ROCm + if: steps.cache-rocm.outputs.cache-hit != 'true' + uses: ./.github/actions/windows-setup-rocm + with: + version: ${{ matrix.ROCM_VERSION }} + + - name: Setup ROCm Environment + run: | + $ErrorActionPreference = "Stop" + + # Activate venv from cache or fresh install + & C:\TheRock\build\.venv\Scripts\Activate.ps1 + + # Expand the devel tree (idempotent; no-op if already done during install) + rocm-sdk init + if ($LASTEXITCODE -ne 0) { throw "rocm-sdk init failed with exit code $LASTEXITCODE" } + + # Get ROCm installation paths using the rocm-sdk CLI tool + $rocmPath = (rocm-sdk path --root) + if (-not $rocmPath) { throw "rocm-sdk path --root returned empty - devel package may not be installed" } + $rocmPath = $rocmPath.Trim() + $cmakePath = (rocm-sdk path --cmake).Trim() + $binPath = (rocm-sdk path --bin).Trim() + write-host "ROCm root: $rocmPath" + write-host "CMake path: $cmakePath" + write-host "Bin path: $binPath" + + echo "HIP_PATH=$rocmPath" >> $env:GITHUB_ENV + echo "CMAKE_PREFIX_PATH=$cmakePath" >> $env:GITHUB_ENV + echo "HIP_DEVICE_LIB_PATH=$rocmPath\lib\llvm\amdgcn\bitcode" >> $env:GITHUB_ENV + echo "HIP_PLATFORM=amd" >> $env:GITHUB_ENV + echo "LLVM_PATH=$rocmPath\lib\llvm" >> $env:GITHUB_ENV + echo "$binPath" >> $env:GITHUB_PATH + + # Keep venv in PATH for subsequent steps + echo "C:\TheRock\build\.venv\Scripts" >> $env:GITHUB_PATH + + - name: Build + run: | + mkdir build + cd build + cmake .. ` + -G "Unix Makefiles" ` + -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_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} + + - name: ccache-clear + uses: ./.github/actions/ccache-clear + with: + key: windows-rocm-${{ matrix.ROCM_VERSION }}-${{ matrix.build }} + + - name: Verify HIP backend was built + run: | + $hipDll = Get-ChildItem -Path build\bin -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 + exit 1 + } + Write-Host "HIP backend artifact found:" + $hipDll | Format-Table FullName, Length -AutoSize + + - name: Determine tag name + id: tag + uses: ./.github/actions/get-tag-name + + - name: Get ROCm short version + run: | + $rocmVersionShort = ('${{ matrix.ROCM_VERSION }}'.Split('.')[0..1] -join '.') + echo "ROCM_VERSION_SHORT=$rocmVersionShort" >> $env:GITHUB_ENV + + - name: Pack artifacts + run: | + cp "LICENSE" "build\bin\" + 7z a -snl llama-bin-win-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.zip .\build\bin\* + + - name: Upload artifacts + uses: actions/upload-artifact@v6 + with: + path: llama-bin-win-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.zip + name: llama-bin-win-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.zip + windows: needs: [check-release] if: ${{ needs.check-release.outputs.should_release == 'true' }} @@ -1168,8 +1295,8 @@ jobs: strategy: matrix: include: - - ROCM_VERSION: "7.2.1" - gpu_targets: "gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1151;gfx1150;gfx1200;gfx1201" + - ROCM_VERSION: "7.14.0" + gpu_targets: "gfx908;gfx90a;gfx942;gfx950;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1152;gfx1200;gfx1201" build: 'x64' steps: @@ -1201,38 +1328,36 @@ jobs: run: | sudo apt install -y build-essential git cmake wget - - name: Setup Legacy ROCm - if: matrix.ROCM_VERSION == '7.2.1' - id: legacy_env - run: | - sudo mkdir --parents --mode=0755 /etc/apt/keyrings - wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | \ - gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null - - sudo tee /etc/apt/sources.list.d/rocm.list << EOF - deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/${{ matrix.ROCM_VERSION }} jammy main - EOF - - sudo tee /etc/apt/preferences.d/rocm-pin-600 << EOF - Package: * - Pin: release o=repo.radeon.com - Pin-Priority: 600 - EOF - - sudo apt update - sudo apt-get install -y libssl-dev rocm-hip-sdk - - - name: Setup TheRock - if: matrix.ROCM_VERSION != '7.2.1' + - name: Setup TheRock with Wheels id: therock_env run: | - wget https://repo.amd.com/rocm/tarball/therock-dist-linux-gfx1151-${{ matrix.ROCM_VERSION }}.tar.gz - mkdir install - tar -xf *.tar.gz -C install - export ROCM_PATH=$(pwd)/install - echo ROCM_PATH=$ROCM_PATH >> $GITHUB_ENV - echo PATH=$PATH:$ROCM_PATH/bin >> $GITHUB_ENV - echo LD_LIBRARY_PATH=$ROCM_PATH/lib:$ROCM_PATH/llvm/lib:$ROCM_PATH/lib/rocprofiler-systems >> $GITHUB_ENV + # Create Python virtual environment + python3 -m venv .venv + source .venv/bin/activate + + # Install ROCm wheels for build + # libraries = HIP runtime and CMake configs needed for linking + # devel = compilers, headers, static libs + python -m pip install --upgrade pip + python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{ matrix.ROCM_VERSION }}" + + # Get ROCm installation paths using the rocm-sdk CLI tool + ROCM_PATH=$(rocm-sdk path --root) + CMAKE_PATH=$(rocm-sdk path --cmake) + BIN_PATH=$(rocm-sdk path --bin) + echo "ROCM_PATH=$ROCM_PATH" + echo "CMAKE_PATH=$CMAKE_PATH" + echo "BIN_PATH=$BIN_PATH" + + # Set environment variables + echo "ROCM_PATH=$ROCM_PATH" >> $GITHUB_ENV + echo "CMAKE_PREFIX_PATH=$CMAKE_PATH" >> $GITHUB_ENV + echo "HIP_PATH=$ROCM_PATH" >> $GITHUB_ENV + echo "PATH=$BIN_PATH:${PATH}" >> $GITHUB_ENV + echo "LD_LIBRARY_PATH=$ROCM_PATH/lib:${LD_LIBRARY_PATH:-}" >> $GITHUB_ENV + + # Keep venv activated for subsequent steps + echo "$(pwd)/.venv/bin" >> $GITHUB_PATH - name: Build with native CMake HIP support id: cmake_build @@ -1276,129 +1401,6 @@ jobs: path: llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz name: llama-bin-ubuntu-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.tar.gz - windows-hip: - needs: [check-release, get-version] - if: ${{ needs.check-release.outputs.should_release == 'true' }} - - runs-on: windows-2022 - - permissions: - actions: write - - env: - HIPSDK_INSTALLER_VERSION: "26.Q1" - - strategy: - matrix: - include: - - name: "radeon" - gpu_targets: "gfx1150;gfx1151;gfx1200;gfx1201;gfx1100;gfx1101;gfx1102;gfx1030;gfx1031;gfx1032" - - steps: - - name: Clone - 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: Grab rocWMMA package - id: grab_rocwmma - run: | - curl -o rocwmma.deb "https://repo.radeon.com/rocm/apt/7.2.1/pool/main/r/rocwmma-dev/rocwmma-dev_2.2.0.70201-81~24.04_amd64.deb" - 7z x rocwmma.deb - 7z x data.tar - - - name: Cache ROCm Installation - id: cache-rocm - uses: actions/cache@v5 - with: - path: C:\Program Files\AMD\ROCm - key: cache-gha-rocm-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ runner.os }} - - - name: ccache - uses: ggml-org/ccache-action@v1.2.21 - with: - key: release-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }} - - - name: Install ROCm - if: steps.cache-rocm.outputs.cache-hit != 'true' - id: depends - run: | - $ErrorActionPreference = "Stop" - write-host "Downloading AMD HIP SDK Installer" - Invoke-WebRequest -Uri "https://download.amd.com/developer/eula/rocm-hub/AMD-Software-PRO-Edition-${{ env.HIPSDK_INSTALLER_VERSION }}-Win11-For-HIP.exe" -OutFile "${env:RUNNER_TEMP}\rocm-install.exe" - write-host "Installing AMD HIP SDK" - $proc = Start-Process "${env:RUNNER_TEMP}\rocm-install.exe" -ArgumentList '-install' -NoNewWindow -PassThru - $completed = $proc.WaitForExit(600000) - if (-not $completed) { - Write-Error "ROCm installation timed out after 10 minutes. Killing the process" - $proc.Kill() - exit 1 - } - if ($proc.ExitCode -ne 0) { - Write-Error "ROCm installation failed with exit code $($proc.ExitCode)" - exit 1 - } - write-host "Completed AMD HIP SDK installation" - - - name: Verify ROCm - id: verify - run: | - # Find and test ROCm installation - $clangPath = Get-ChildItem 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | Select-Object -First 1 - if (-not $clangPath) { - Write-Error "ROCm installation not found" - exit 1 - } - & $clangPath.FullName --version - - - name: Build - id: cmake_build - run: | - $env:HIP_PATH=$(Resolve-Path 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | split-path | split-path) - $env:CMAKE_PREFIX_PATH="${env:HIP_PATH}" - cmake -G "Unix Makefiles" -B build -S . ` - -DCMAKE_C_COMPILER="${env:HIP_PATH}\bin\clang.exe" ` - -DCMAKE_CXX_COMPILER="${env:HIP_PATH}\bin\clang++.exe" ` - -DCMAKE_CXX_FLAGS="-I$($PWD.Path.Replace('\', '/'))/opt/rocm-7.2.1/include/ -Wno-ignored-attributes -Wno-nested-anon-types" ` - -DCMAKE_BUILD_TYPE=Release ` - -DGGML_BACKEND_DL=ON ` - -DGGML_NATIVE=OFF ` - -DGGML_CPU=OFF ` - -DGPU_TARGETS="${{ matrix.gpu_targets }}" ` - -DGGML_HIP=ON ` - -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} ` - -DLLAMA_BUILD_BORINGSSL=ON - cmake --build build --target ggml-hip -j ${env:NUMBER_OF_PROCESSORS} - md "build\bin\rocblas\library\" - md "build\bin\hipblaslt\library" - cp "${env:HIP_PATH}\bin\libhipblas.dll" "build\bin\" - cp "${env:HIP_PATH}\bin\libhipblaslt.dll" "build\bin\" - cp "${env:HIP_PATH}\bin\rocblas.dll" "build\bin\" - cp "${env:HIP_PATH}\bin\rocblas\library\*" "build\bin\rocblas\library\" - cp "${env:HIP_PATH}\bin\hipblaslt\library\*" "build\bin\hipblaslt\library\" - - - name: ccache-clear - uses: ./.github/actions/ccache-clear - with: - key: release-windows-2022-x64-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ matrix.name }} - - - name: Pack artifacts - id: pack_artifacts - run: | - 7z a -snl llama-bin-win-hip-${{ matrix.name }}-x64.zip .\build\bin\* - - - name: Upload artifacts - uses: actions/upload-artifact@v6 - with: - path: llama-bin-win-hip-${{ matrix.name }}-x64.zip - name: llama-bin-win-hip-${{ matrix.name }}-x64.zip - ios-xcode: needs: [check-release, get-version] if: ${{ needs.check-release.outputs.should_release == 'true' }} @@ -1572,7 +1574,7 @@ jobs: - windows-cpu - windows-cuda #- windows-sycl - - windows-hip + - windows-rocm - windows-openvino - ubuntu-22-rocm - ubuntu-cpu @@ -1684,7 +1686,7 @@ jobs: - [Ubuntu s390x (CPU)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-s390x.tar.gz) - [Ubuntu x64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-x64.tar.gz) - [Ubuntu arm64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-vulkan-arm64.tar.gz) - - [Ubuntu x64 (ROCm 7.2)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-7.2-x64.tar.gz) + - [Ubuntu x64 (ROCm 7.14)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-rocm-7.14-x64.tar.gz) - [Ubuntu x64 (OpenVINO)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-openvino-${{ needs.ubuntu-24-openvino.outputs.openvino_version }}-x64.tar.gz) - [Ubuntu x64 (SYCL FP32)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-fp32-x64.tar.gz) - [Ubuntu x64 (SYCL FP16)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-ubuntu-sycl-fp16-x64.tar.gz) @@ -1702,7 +1704,7 @@ jobs: - [Windows x64 (Vulkan)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-vulkan-x64.zip) - [Windows x64 (OpenVINO)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-openvino-${{ needs.windows-openvino.outputs.openvino_version }}-x64.zip) - [Windows x64 (SYCL)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-sycl-x64.zip) - - [Windows x64 (HIP)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-hip-radeon-x64.zip) + - [Windows x64 (ROCm 7.14)](https://github.com/ggml-org/llama.cpp/releases/download/${{ steps.tag.outputs.name }}/llama-${{ steps.tag.outputs.name }}-bin-win-rocm-7.14-x64.zip) **openEuler:** - [DISABLED](https://github.com/ggml-org/llama.cpp/pull/23705) diff --git a/.github/workflows/server-sanitize.yml b/.github/workflows/server-sanitize.yml index 0eeefdf88469..5d696282c709 100644 --- a/.github/workflows/server-sanitize.yml +++ b/.github/workflows/server-sanitize.yml @@ -110,7 +110,7 @@ jobs: source .venv/bin/activate cd tools/server/tests export ${{ matrix.extra_args }} - pytest -v -x -m "not slow" + ./tests.sh - name: Slow tests id: server_integration_tests_slow @@ -119,4 +119,4 @@ jobs: source .venv/bin/activate cd tools/server/tests export ${{ matrix.extra_args }} - SLOW_TESTS=1 pytest -v -x + SLOW_TESTS=1 ./tests.sh diff --git a/.github/workflows/server-self-hosted.yml b/.github/workflows/server-self-hosted.yml index 249f389ff3f2..675ddbaaa580 100644 --- a/.github/workflows/server-self-hosted.yml +++ b/.github/workflows/server-self-hosted.yml @@ -72,7 +72,7 @@ jobs: run: | cd tools/server/tests source venv/bin/activate - pytest -v -x -m "not slow" + ./tests.sh - name: Tests (GPUx1, backend-sampling) id: server_integration_tests_backend_sampling @@ -81,7 +81,7 @@ jobs: cd tools/server/tests source venv/bin/activate export LLAMA_ARG_BACKEND_SAMPLING=1 - pytest -v -x -m "not slow" + ./tests.sh - name: Tests (GPUx2) id: server_integration_tests_gpu2 @@ -90,7 +90,7 @@ jobs: cd tools/server/tests source venv/bin/activate export GGML_METAL_DEVICES=2 - pytest -v -x -m "not slow" + ./tests.sh - name: Tests (GPUx2, backend-sampling) id: server_integration_tests_gpu2_backend_sampling @@ -99,7 +99,7 @@ jobs: cd tools/server/tests source venv/bin/activate export GGML_METAL_DEVICES=2 LLAMA_ARG_BACKEND_SAMPLING=1 - pytest -v -x -m "not slow" + ./tests.sh server-cuda: runs-on: [self-hosted, llama-server, Linux, NVIDIA] @@ -132,7 +132,7 @@ jobs: run: | cd tools/server/tests source venv/bin/activate - pytest -v -x -m "not slow" + ./tests.sh - name: Tests (GPUx1, backend-sampling) id: server_integration_tests_backend_sampling @@ -141,7 +141,7 @@ jobs: cd tools/server/tests source venv/bin/activate export LLAMA_ARG_BACKEND_SAMPLING=1 - pytest -v -x -m "not slow" + ./tests.sh - name: Tests (GPUx2) id: server_integration_tests_gpu2 @@ -150,7 +150,7 @@ jobs: cd tools/server/tests source venv/bin/activate export GGML_CUDA_DEVICES=2 - pytest -v -x -m "not slow" + ./tests.sh - name: Tests (GPUx2, backend-sampling) id: server_integration_tests_gpu2_backend_sampling @@ -159,7 +159,7 @@ jobs: cd tools/server/tests source venv/bin/activate export GGML_CUDA_DEVICES=2 LLAMA_ARG_BACKEND_SAMPLING=1 - pytest -v -x -m "not slow" + ./tests.sh server-kleidiai: runs-on: ah-ubuntu_22_04-c8g_8x @@ -219,4 +219,4 @@ jobs: run: | cd tools/server/tests source venv/bin/activate - pytest -v -x -m "not slow" + ./tests.sh diff --git a/.github/workflows/server.yml b/.github/workflows/server.yml index 5a02cc15ad5e..9fb4b4ba102f 100644 --- a/.github/workflows/server.yml +++ b/.github/workflows/server.yml @@ -104,21 +104,21 @@ jobs: id: server_integration_tests run: | cd tools/server/tests - pytest -v -x -m "not slow" + ./tests.sh - name: Slow tests id: server_integration_tests_slow if: ${{ github.event.schedule || github.event.inputs.slow_tests == 'true' }} run: | cd tools/server/tests - SLOW_TESTS=1 pytest -v -x + SLOW_TESTS=1 ./tests.sh - name: Tests (Backend sampling) id: server_integration_tests_backend_sampling run: | cd tools/server/tests export LLAMA_ARG_BACKEND_SAMPLING=1 - pytest -v -x -m "not slow" + ./tests.sh - name: Slow tests (Backend sampling) id: server_integration_tests_slow_backend_sampling @@ -126,7 +126,7 @@ jobs: run: | cd tools/server/tests export LLAMA_ARG_BACKEND_SAMPLING=1 - SLOW_TESTS=1 pytest -v -x + SLOW_TESTS=1 ./tests.sh windows: runs-on: windows-2025 @@ -167,15 +167,17 @@ jobs: - name: Tests id: server_integration_tests + shell: bash run: | cd tools/server/tests - $env:PYTHONIOENCODING = ":replace" - pytest -v -x -m "not slow" + export PYTHONIOENCODING=":replace" + ./tests.sh - name: Slow tests id: server_integration_tests_slow if: ${{ github.event.schedule || github.event.inputs.slow_tests == 'true' }} + shell: bash run: | cd tools/server/tests - $env:SLOW_TESTS = "1" - pytest -v -x + export SLOW_TESTS="1" + ./tests.sh diff --git a/ci/run.sh b/ci/run.sh index f6c7eb0d5c70..8046df255154 100755 --- a/ci/run.sh +++ b/ci/run.sh @@ -49,6 +49,14 @@ mkdir -p "$2" OUT=$(realpath "$1") MNT=$(realpath "$2") +# gpu-rocm self-hosted runner can't upload logs to blob; keep each run's logs in +# their own dir keyed by the GitHub run id so an Actions run URL maps to its logs. +if [ -n "${GG_BUILD_ROCM}" ] && [ -n "${GITHUB_RUN_ID}" ]; then + OUT="$OUT/run-${GITHUB_RUN_ID}-${GITHUB_RUN_ATTEMPT:-1}" + mkdir -p "$OUT" + echo "ci results dir: $OUT" +fi + rm -f $OUT/*.log rm -f $OUT/*.exit rm -f $OUT/*.md diff --git a/common/arg.cpp b/common/arg.cpp index c37d5cd0aa4b..cb314eee7044 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2605,14 +2605,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ).set_env("LLAMA_ARG_DIO")); add_opt(common_arg( {"-lm", "--load-mode"}, "MODE", - "model loading mode (default: mmap)\n" + "model loading mode (default: auto)\n" + "- auto: mmap, unless a device does not support it\n" "- none: no special loading mode\n" "- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)\n" "- mlock: force system to keep model in RAM rather than swapping or compressing\n" "- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing\n" "- dio: use DirectIO if available\n", [](common_params & params, const std::string & value) { - /**/ if (value == "none") { params.load_mode = LLAMA_LOAD_MODE_NONE; } + /**/ if (value == "auto") { params.load_mode = LLAMA_LOAD_MODE_AUTO; } + else if (value == "none") { params.load_mode = LLAMA_LOAD_MODE_NONE; } else if (value == "mmap") { params.load_mode = LLAMA_LOAD_MODE_MMAP; } else if (value == "mlock") { params.load_mode = LLAMA_LOAD_MODE_MLOCK; } else if (value == "mmap+mlock") { params.load_mode = LLAMA_LOAD_MODE_MMAP_MLOCK; } diff --git a/common/chat.cpp b/common/chat.cpp index 6cbf23b50508..01053ddde804 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -1166,6 +1166,16 @@ static common_chat_params common_chat_params_init_qwen3_coder(const common_chat_ data.prompt += data.generation_prompt; } + std::vector tool_call_starts = { "" }; + + // Match complete opener for Qwen3-Coder models that occasionally omit the + // starting . The model may hallucinate a tool name, but it is preferable over + // constraining on + foreach_function(inputs.tools, [&](const json & tool) { + const std::string name = tool.at("function").at("name"); + tool_call_starts.push_back(""); + }); + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { auto generation_prompt = p.literal(GEN_PREFIX); @@ -1238,7 +1248,7 @@ static common_chat_params common_chat_params_init_qwen3_coder(const common_chat_ auto tool_calls = p.trigger_rule("tool-call-root", p.repeat(calls, min_calls, 1)); return generation_prompt + - (reasoning << p.content(p.until_one_of({ "", "" }, - // Trigger on " fit_params_target = std::vector(llama_max_devices(), 1024 * 1024*1024); 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_MMAP; // how to load the model + enum llama_load_mode load_mode = LLAMA_LOAD_MODE_AUTO; // how to load the model common_cpu_params cpuparams; common_cpu_params cpuparams_batch; diff --git a/common/peg-parser.cpp b/common/peg-parser.cpp index ef290ed7c057..4a4be7cf789f 100644 --- a/common/peg-parser.cpp +++ b/common/peg-parser.cpp @@ -570,23 +570,34 @@ struct parser_executor { } static common_peg_parse_result handle_escape_sequence(common_peg_parse_context & ctx, size_t start, size_t & pos, const char delimiter) { + auto save = pos; + ++pos; // consume '\' if (pos >= ctx.input.size()) { if (!ctx.is_lenient()) { return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start); } + pos = save; // suppress unmatched '\' return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start, pos); } char c = ctx.input[pos]; + if (c == delimiter || c == '\\' || c == '/' || c == 'b' || c == 'f' || c == 'n' || c == 'r' || c == 't') { ++pos; return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_SUCCESS, start, pos); - } else if (c == 'u') { - return handle_unicode_escape(ctx, start, pos); - } else { - return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start); } + + if (c == 'u') { + auto result = handle_unicode_escape(ctx, start, pos); + if (result.need_more_input()) { + pos = save; // suppress incomplete sequence + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_NEED_MORE_INPUT, start, pos); + } + return result; + } + + return common_peg_parse_result(COMMON_PEG_PARSE_RESULT_FAIL, start); } static common_peg_parse_result handle_unicode_escape(common_peg_parse_context & ctx, size_t start, size_t & pos) { diff --git a/conversion/__init__.py b/conversion/__init__.py index c7d8046c495f..695289b73ace 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -214,6 +214,7 @@ "Qwen3MoeForCausalLM": "qwen", "Qwen3NextForCausalLM": "qwen", "Qwen3OmniMoeForConditionalGeneration": "qwen3vl", + "PocketTTSModel": "pockettts", "Qwen3TTSForConditionalGeneration": "qwen3tts", "Qwen3VLForConditionalGeneration": "qwen3vl", "Qwen3VLMoeForConditionalGeneration": "qwen3vl", @@ -310,6 +311,7 @@ "Qwen2_5_VLForConditionalGeneration": "qwenvl", "Qwen3ASRForConditionalGeneration": "qwen3vl", "Qwen3OmniMoeForConditionalGeneration": "qwen3vl", + "PocketTTSModel": "pockettts", "Qwen3TTSForConditionalGeneration": "qwen3tts", "Qwen3VLForConditionalGeneration": "qwen3vl", "Qwen3VLMoeForConditionalGeneration": "qwen3vl", diff --git a/conversion/base.py b/conversion/base.py index a7cd3fd904aa..718d5394495e 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -58,6 +58,11 @@ AnyModel = TypeVar("AnyModel", bound="type[ModelBase]") +# for checkpoints that ship no config.json, we will try to provide a synthetic one +HparamsMatcher = Callable[[Path], bool] +HparamsLoader = Callable[[Path], dict[str, Any]] + + class SentencePieceTokenTypes(IntEnum): NORMAL = 1 UNKNOWN = 2 @@ -77,6 +82,7 @@ class ModelBase: ModelType.TEXT: {}, ModelType.MMPROJ: {}, } + _hparams_loaders: list[tuple[HparamsMatcher, HparamsLoader]] = [] dir_model: Path ftype: gguf.LlamaFileType @@ -823,7 +829,7 @@ def prepare_tensors(self): elif any(str(v.get("quant_algo")).endswith("NVFP4") for v in quant_layers.values() if isinstance(v, dict)): quant_algo = "NVFP4" - self._is_nvfp4 = quant_algo == "NVFP4" + self._is_nvfp4 = quant_algo in ("NVFP4", "W4A16_NVFP4") self._is_mxfp4 = quant_method == "mxfp4" # NVFP4 weights are repacked and written directly to gguf_writer. @@ -1040,6 +1046,24 @@ def get_model_part_names(dir_model: Path, prefix: str, suffix: str) -> list[str] return part_names + @staticmethod + def load_hparams_guess(dir_model: Path) -> dict[str, Any] | None: + # some models ship no config.json, will try to guess them + from conversion import load_all_models + load_all_models() + + for matcher, loader in ModelBase._hparams_loaders: + if matcher(dir_model): + return loader(dir_model) + return None + + @classmethod + def register_hparams_loader(cls, matcher: HparamsMatcher) -> Callable[[HparamsLoader], HparamsLoader]: + def inner(loader: HparamsLoader) -> HparamsLoader: + cls._hparams_loaders.append((matcher, loader)) + return loader + return inner + @staticmethod def load_hparams(dir_model: Path, is_mistral_format: bool): if is_mistral_format: @@ -1053,6 +1077,10 @@ def load_hparams(dir_model: Path, is_mistral_format: bool): config = AutoConfig.from_pretrained(dir_model, trust_remote_code=False).to_dict() except Exception as e: logger.warning(f"Failed to load model config from {dir_model}: {e}") + if not (dir_model / "config.json").is_file(): + config = ModelBase.load_hparams_guess(dir_model) + if config is not None: + return config logger.warning("Trying to load config.json instead") with open(dir_model / "config.json", "r", encoding="utf-8") as f: config = json.load(f) diff --git a/conversion/nemotron.py b/conversion/nemotron.py index e5075020c1e3..c46cec143866 100644 --- a/conversion/nemotron.py +++ b/conversion/nemotron.py @@ -275,10 +275,18 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca return None elif cls.mtp_only: # --mtp: export the MTP head plus the tensors it shares with the target model + # Include lm_head scale sidecars so NVFP4 packing sees them. keep = name in ( "backbone.embeddings.weight", "backbone.norm_f.weight", "lm_head.weight", + "lm_head.weight_scale", + "lm_head.weight_scale_2", + "lm_head.weight_scale_inv", + "lm_head.input_scale", + "lm_head.input_global_scale", + "lm_head.weight_global_scale", + "lm_head.weight_packed", ) if not keep: return None diff --git a/conversion/pockettts.py b/conversion/pockettts.py new file mode 100644 index 000000000000..62ecb5acde74 --- /dev/null +++ b/conversion/pockettts.py @@ -0,0 +1,378 @@ +from __future__ import annotations + +import re +from pathlib import Path +from typing import Any, Iterable, TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from torch import Tensor + +from .base import ModelBase, MmprojModel, SentencePieceTokenTypes, TextModel, gguf, logger + +# Pocket TTS is a CALM: the backbone conditions a flow-matching decoder that generates one +# continuous 32-d latent per frame. There is no codebook in this model. +# The checkpoint ships no config.json, hparams come from _load_hparams() below. +# +# Tricks being used to support this model via existing llama.cpp code paths: +# - bos_before_voice and bos_emb are learned input vectors, not tokens +# they are appended to the embedding table as extra tokens, to be looked up like any other row +# - bos_emb lives in latent space, so input_linear is folded into it here +# - the backbone has no lm_head, the embedding table is reused as output for the unused logits +# +# pipeline stage mapping: +# mimi encoder + speaker_proj --> mapped to normal mtmd audio encoder +# flow_lm.transformer --> mapped to normal libllama text model (autoregressive) +# flow_lm.flow_net + out_eos --> MTMD_GEN_PROCESS_TYPE_GEN_CODE +# mimi decoder --> MTMD_GEN_PROCESS_TYPE_GEN_WAV + +# indices into mimi.encoder.model / mimi.decoder.model for stage i, see SEANetEncoder/SEANetDecoder +_ENC_RES_IDX = lambda i: 1 + 3 * i # noqa: E731 +_ENC_SCALE_IDX = lambda i: 3 + 3 * i # noqa: E731 +_DEC_SCALE_IDX = lambda i: 2 + 3 * i # noqa: E731 +_DEC_RES_IDX = lambda i: 3 + 3 * i # noqa: E731 + +_N_SEANET_STAGES = 3 +_SAMPLE_RATE = 24000 + + +def _tensor_shapes(dir_model: Path) -> dict[str, tuple[int, ...]]: + part_names = ModelBase.get_model_part_names(dir_model, "model", ".safetensors") + if len(part_names) != 1: + return {} + with gguf.utility.SafetensorsLocal(dir_model / part_names[0]) as part: + return {name: tuple(part[name].shape) for name in part.keys()} + + +@ModelBase.register_hparams_loader(lambda dir_model: "flow_lm.bos_emb" in _tensor_shapes(dir_model)) +def _load_hparams(dir_model: Path) -> dict[str, Any]: + logger.info("gguf: detected pocket-tts checkpoint, deriving hparams from tensor shapes") + shapes = _tensor_shapes(dir_model) + n_vocab, n_embd = shapes["flow_lm.conditioner.embed.weight"] + n_layer = sum(1 for name in shapes if re.fullmatch(r"flow_lm\.transformer\.layers\.\d+\.norm1\.weight", name)) + n_layer_a = sum(1 for name in shapes if re.fullmatch(r"mimi\.encoder_transformer\.transformer\.layers\.\d+\.norm1\.weight", name)) + n_embd_a = shapes["mimi.encoder_transformer.transformer.layers.0.norm1.weight"][0] + return { + "architectures": ["PocketTTSModel"], + "model_type": "pockettts", + "num_hidden_layers": n_layer, + "hidden_size": n_embd, + "intermediate_size": shapes["flow_lm.transformer.layers.0.linear1.weight"][0], + # the transformer is fully causal with no context limit, this only bounds the KV cache + "max_position_embeddings": 4096, + # not in the checkpoint, but every released variant uses head_dim 64 + "num_attention_heads": n_embd // 64, + # extra rows for the learned input vectors, see _embd_table() + "vocab_size": n_vocab + (2 if "flow_lm.bos_before_voice" in shapes else 1), + "rope_theta": 10000.0, + "layer_norm_eps": 1e-5, + "audio_config": { + "num_hidden_layers": n_layer_a, + "hidden_size": n_embd_a, + "intermediate_size": shapes["mimi.encoder_transformer.transformer.layers.0.linear1.weight"][0], + "num_attention_heads": n_embd_a // 64, + }, + } + + +@ModelBase.register("PocketTTSModel") +class PocketTTSModel(TextModel): + model_arch = gguf.MODEL_ARCH.POCKETTTS + + _LAYER_TENSOR_MAP = { + "norm1": gguf.MODEL_TENSOR.ATTN_NORM, + "norm2": gguf.MODEL_TENSOR.FFN_NORM, + "self_attn.out_proj": gguf.MODEL_TENSOR.ATTN_OUT, + "linear1": gguf.MODEL_TENSOR.FFN_UP, + "linear2": gguf.MODEL_TENSOR.FFN_DOWN, + } + + def set_vocab(self): + # this is a unigram sentencepiece model, llama.cpp's SPM tokenizer cannot do + # unigram segmentation, so use the UGM tokenizer instead + from sentencepiece import sentencepiece_model_pb2 as model + + proto = model.ModelProto() # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute] + proto.ParseFromString(open(self.dir_model / "tokenizer.model", "rb").read()) + assert proto.trainer_spec.model_type == 1, "expected a unigram tokenizer" + + tokens, scores, toktypes = self._create_vocab_sentencepiece() + + # the last rows of the embedding table are not sentencepiece pieces + extra = self._extra_tokens() + for i, name in enumerate(extra): + tokens[len(tokens) - len(extra) + i] = name.encode("utf-8") + toktypes[len(tokens) - len(extra) + i] = SentencePieceTokenTypes.CONTROL + scores[len(tokens) - len(extra) + i] = -1000.0 + + self.gguf_writer.add_tokenizer_model("t5") + self.gguf_writer.add_tokenizer_pre("default") + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_scores(scores) + self.gguf_writer.add_token_types(toktypes) + self.gguf_writer.add_add_space_prefix(proto.normalizer_spec.add_dummy_prefix) + self.gguf_writer.add_remove_extra_whitespaces(proto.normalizer_spec.remove_extra_whitespaces) + if proto.normalizer_spec.precompiled_charsmap: + self.gguf_writer.add_precompiled_charsmap(proto.normalizer_spec.precompiled_charsmap) + self.gguf_writer.add_add_bos_token(False) + self.gguf_writer.add_add_eos_token(False) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if not name.startswith("flow_lm."): + return # mimi and the flow net go to the mmproj + + if name == "flow_lm.conditioner.embed.weight": + yield (self.format_tensor_name(gguf.MODEL_TENSOR.TOKEN_EMBD), self._embd_table(data_torch)) + return + + if name.startswith("flow_lm.out_norm."): + suffix = "." + name.rsplit(".", 1)[1] + yield (self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT_NORM, suffix=suffix), data_torch) + return + + if name.startswith("flow_lm.transformer.layers."): + assert bid is not None + key_with_suffix = name.split(f"layers.{bid}.", 1)[1] + key, suffix = key_with_suffix.rsplit(".", 1) + + if key == "self_attn.in_proj": + q, k, v = data_torch.chunk(3, dim=0) + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_Q, bid), q) + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K, bid), k) + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V, bid), v) + return + + tensor = self._LAYER_TENSOR_MAP.get(key) + if tensor is not None: + yield (self.format_tensor_name(tensor, bid, suffix="." + suffix), data_torch) + return + + return + + def _extra_tokens(self) -> list[str]: + # the conditioner's padding row, then the learned vectors appended by _embd_table(). + # bos_before_voice only exists when the pack sets insert_bos_before_voice + names = ["<|pad|>"] + if "flow_lm.bos_before_voice" in self.model_tensors: + names.append("<|bos_before_voice|>") + names.append("<|audio_bos|>") + return names + + def _embd_table(self, embed: Tensor) -> Tensor: + rows = [embed] + if "flow_lm.bos_before_voice" in self.model_tensors: + rows.append(self.model_tensors["flow_lm.bos_before_voice"]().reshape(1, -1).to(embed.dtype)) + + # bos_emb is a latent, it only enters the backbone through input_linear + bos_emb = self.model_tensors["flow_lm.bos_emb"]() + input_linear = self.model_tensors["flow_lm.input_linear.weight"]() + audio_bos = torch.nn.functional.linear(bos_emb.float(), input_linear.float()).reshape(1, -1) + rows.append(audio_bos.to(embed.dtype)) + + return torch.cat(rows, dim=0) + + +@ModelBase.register("PocketTTSModel") +class PocketTTSMmprojModel(MmprojModel): + has_audio_encoder = True + has_vision_encoder = False + + _MIMI_TFM_MAP = { + "norm1": (gguf.MODEL_TENSOR.A_ENC_INPUT_NORM, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_NORM), + "norm2": (gguf.MODEL_TENSOR.A_ENC_OUTPUT_NORM, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_NORM), + "self_attn.out_proj": (gguf.MODEL_TENSOR.A_ENC_OUTPUT, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_OUT), + "linear1": (gguf.MODEL_TENSOR.A_ENC_FFN_UP, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_UP), + "linear2": (gguf.MODEL_TENSOR.A_ENC_FFN_DOWN, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_DOWN), + "layer_scale_1.scale": (gguf.MODEL_TENSOR.A_ENC_ATTN_SCALE, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_SCALE), + "layer_scale_2.scale": (gguf.MODEL_TENSOR.A_ENC_FFN_SCALE_LS, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_SCALE), + } + _MIMI_TFM_QKV = ( + (gguf.MODEL_TENSOR.A_ENC_ATTN_Q, gguf.MODEL_TENSOR.A_ENC_ATTN_K, gguf.MODEL_TENSOR.A_ENC_ATTN_V), + (gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_Q, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_K, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_V), + ) + + def set_gguf_parameters(self): + self.gguf_writer.add_file_type(self.ftype) + assert self.hparams_audio is not None + + # voice-prompt encoder: mimi encoder + speaker_proj + self.gguf_writer.add_clip_has_audio_encoder(True) + # note: the 24kHz sample rate is hardcoded on the clip.cpp side, like the other audio models + self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_SPKENC) + self.gguf_writer.add_audio_projection_dim(self.n_embd_text) + self.gguf_writer.add_audio_block_count(self.hparams_audio["num_hidden_layers"]) + self.gguf_writer.add_audio_embedding_length(self.hparams_audio["hidden_size"]) + self.gguf_writer.add_audio_feed_forward_length(self.hparams_audio["intermediate_size"]) + self.gguf_writer.add_audio_head_count(self.hparams_audio["num_attention_heads"]) + self.gguf_writer.add_audio_attention_layernorm_eps(1e-5) + # mimi convolves the waveform directly, it is passed around as a 1-row "mel" + self.gguf_writer.add_audio_num_mel_bins(1) + + # generation: flow-matching decoder + mimi decoder + # the SEANet and flow net hparams are constant across the family, clip.cpp holds them + self.gguf_writer.add_clip_has_gen_audio_encoder(True) + self.gguf_writer.add_clip_gen_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_GEN) + self.gguf_writer.add_gen_audio_projection_dim(self.n_embd_text) + self.gguf_writer.add_gen_audio_embedding_length(self.hparams_audio["hidden_size"]) + self.gguf_writer.add_gen_audio_feed_forward_length(self.hparams_audio["intermediate_size"]) + self.gguf_writer.add_gen_audio_block_count(self.hparams_audio["num_hidden_layers"]) + self.gguf_writer.add_gen_audio_head_count(self.hparams_audio["num_attention_heads"]) + self.gguf_writer.add_gen_audio_attention_layernorm_eps(1e-5) + + self.gguf_writer.add_gen_audio_model_variant(self.dir_model.name) + + def tensor_force_quant(self, name, new_name, bid, n_dims): + del name, bid, n_dims + # conv1d/conv1d_dw kernels must be F16, ggml_conv_1d(_dw) has no BF16 path + if ".seanet." in new_name or new_name in ("a.downsample.conv.weight", "a.gen.wav.upsample.weight"): + return gguf.GGMLQuantizationType.F16 + return False + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # the block index of the mimi transformers is parsed here, not by the base class + T = gguf.MODEL_TENSOR + + if name in ("flow_lm.bos_emb", "flow_lm.bos_before_voice", "flow_lm.conditioner.embed.weight"): + return # folded into the backbone embedding table + if name.startswith("flow_lm.transformer.") or name.startswith("flow_lm.out_norm."): + return # backbone + + if name == "flow_lm.speaker_proj_weight": + yield (self.format_tensor_name(T.A_ENC_SPEAKER_PROJ), data_torch) + return + if name == "flow_lm.input_linear.weight": + yield (self.format_tensor_name(T.A_GEN_INPUT_LINEAR), data_torch) + return + if name == "flow_lm.emb_mean": + yield (self.format_tensor_name(T.A_GEN_EMB_MEAN, suffix=""), data_torch) + return + if name == "flow_lm.emb_std": + yield (self.format_tensor_name(T.A_GEN_EMB_STD, suffix=""), data_torch) + return + if name.startswith("flow_lm.out_eos."): + suffix = "." + name.rsplit(".", 1)[1] + yield (self.format_tensor_name(T.A_GEN_OUT_EOS, suffix=suffix), data_torch) + return + + if name.startswith("flow_lm.flow_net."): + yield from self._flow_net_tensor(name, data_torch) + return + + if name == "mimi.downsample.conv.conv.weight": + yield (self.format_tensor_name(T.A_ENC_DOWNSAMPLE_CONV), data_torch) + return + if name == "mimi.upsample.convtr.convtr.weight": + yield (self.format_tensor_name(T.A_GEN_WAV_UPSAMPLE), data_torch) + return + if name == "mimi.quantizer.output_proj.weight": + yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_OUT), data_torch.squeeze(-1)) + return + + if "_transformer.transformer.layers." in name: + yield from self._mimi_tfm_tensor(name, data_torch) + return + + if name.startswith("mimi.encoder.model.") or name.startswith("mimi.decoder.model."): + yield from self._seanet_tensor(name, data_torch) + return + + return + + def _flow_net_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]: + T = gguf.MODEL_TENSOR + key = name.split("flow_lm.flow_net.", 1)[1] + suffix = "." + key.rsplit(".", 1)[1] + + simple = { + "input_proj": T.A_GEN_FLOW_INPUT_PROJ, + "cond_embed": T.A_GEN_FLOW_COND_EMBD, + "final_layer.linear": T.A_GEN_FLOW_FINAL_PROJ, + "final_layer.adaLN_modulation.1": T.A_GEN_FLOW_FINAL_ADA, + } + tensor = simple.get(key.rsplit(".", 1)[0]) + if tensor is not None: + yield (self.format_tensor_name(tensor, suffix=suffix), data_torch) + return + + if key.startswith("time_embed."): + bid = int(key.split(".")[1]) + rest = key.split(f"time_embed.{bid}.", 1)[1] + time_map = { + "freqs": (T.A_GEN_FLOW_TIME_FREQS, ""), + "mlp.0": (T.A_GEN_FLOW_TIME_UP, suffix), + "mlp.2": (T.A_GEN_FLOW_TIME_DOWN, suffix), + "mlp.3.alpha": (T.A_GEN_FLOW_TIME_NORM, ""), + } + entry = time_map.get(rest) or time_map.get(rest.rsplit(".", 1)[0]) + if entry is not None: + yield (self.format_tensor_name(entry[0], bid, suffix=entry[1]), data_torch) + return + + if key.startswith("res_blocks."): + bid = int(key.split(".")[1]) + rest = key.split(f"res_blocks.{bid}.", 1)[1].rsplit(".", 1)[0] + blk_map = { + "in_ln": T.A_GEN_FLOW_BLK_NORM, + "mlp.0": T.A_GEN_FLOW_BLK_UP, + "mlp.2": T.A_GEN_FLOW_BLK_DOWN, + "adaLN_modulation.1": T.A_GEN_FLOW_BLK_ADA, + } + tensor = blk_map.get(rest) + if tensor is not None: + yield (self.format_tensor_name(tensor, bid, suffix=suffix), data_torch) + return + + def _mimi_tfm_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]: + is_decoder = name.startswith("mimi.decoder_transformer.") + bid = int(name.split("_transformer.transformer.layers.", 1)[1].split(".")[0]) + key_with_suffix = name.split(f".layers.{bid}.", 1)[1] + + if key_with_suffix == "self_attn.in_proj.weight": + q, k, v = data_torch.chunk(3, dim=0) + names = self._MIMI_TFM_QKV[1 if is_decoder else 0] + for tensor, part in zip(names, (q, k, v)): + yield (self.format_tensor_name(tensor, bid), part) + return + + key, suffix = key_with_suffix.rsplit(".", 1) + entry = self._MIMI_TFM_MAP.get(key) or self._MIMI_TFM_MAP.get(key_with_suffix) + if entry is None: + return + tensor = entry[1 if is_decoder else 0] + suffix = ".weight" if key_with_suffix.endswith(".scale") else "." + suffix + yield (self.format_tensor_name(tensor, bid, suffix=suffix), data_torch) + + def _seanet_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]: + T = gguf.MODEL_TENSOR + is_decoder = name.startswith("mimi.decoder.") + idx = int(name.split(".model.", 1)[1].split(".")[0]) + suffix = "." + name.rsplit(".", 1)[1] + + conv_in, conv_out, res1, res2, scale = ( + (T.A_GEN_WAV_SEANET_CONV_IN, T.A_GEN_WAV_SEANET_CONV_OUT, T.A_GEN_WAV_SEANET_RES_CONV1, + T.A_GEN_WAV_SEANET_RES_CONV2, T.A_GEN_WAV_SEANET_SCALE_CONV) + if is_decoder else + (T.A_ENC_SEANET_CONV_IN, T.A_ENC_SEANET_CONV_OUT, T.A_ENC_SEANET_RES_CONV1, + T.A_ENC_SEANET_RES_CONV2, T.A_ENC_SEANET_SCALE_CONV) + ) + + if idx == 0: + yield (self.format_tensor_name(conv_in, suffix=suffix), data_torch) + return + if idx == 3 * _N_SEANET_STAGES + 2: + yield (self.format_tensor_name(conv_out, suffix=suffix), data_torch) + return + + for stage in range(_N_SEANET_STAGES): + res_idx = _DEC_RES_IDX(stage) if is_decoder else _ENC_RES_IDX(stage) + scale_idx = _DEC_SCALE_IDX(stage) if is_decoder else _ENC_SCALE_IDX(stage) + if idx == scale_idx: + yield (self.format_tensor_name(scale, stage, suffix=suffix), data_torch) + return + if idx == res_idx: + # block.1 is the dilated conv, block.3 the pointwise one (0 and 2 are ELU) + inner = int(name.split(".block.", 1)[1].split(".")[0]) + tensor = res1 if inner == 1 else res2 + yield (self.format_tensor_name(tensor, stage, suffix=suffix), data_torch) + return diff --git a/conversion/qwen.py b/conversion/qwen.py index b4ae528bf2d4..ead435455dd6 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -647,10 +647,13 @@ def set_vocab(self): # own tokenizer logic, not the Qwen default). from . import get_model_class with open(self.target_model_dir / "config.json", "r", encoding="utf-8") as f: - target_arch = json.load(f)["architectures"][0] + target_hparams = json.load(f) + target_arch = target_hparams["architectures"][0] target_cls = get_model_class(target_arch) if target_cls is not type(self): + if target_arch == "NemotronHForCausalLM": + setattr(self, "is_moe", "num_experts_per_tok" in target_hparams) target_cls.set_vocab(self) # ty: ignore[unresolved-attribute] else: super().set_vocab() @@ -688,6 +691,12 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca name = "model." + name return super().filter_tensors((name, gen)) + 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 + + yield from super().modify_tensors(data_torch, name, bid) + @ModelBase.register("Qwen3DSparkModel") class DSparkModel(DFlashModel): diff --git a/examples/model-conversion/requirements.txt b/examples/model-conversion/requirements.txt index 229b2ec75b75..d2cd357ec95c 100644 --- a/examples/model-conversion/requirements.txt +++ b/examples/model-conversion/requirements.txt @@ -1,6 +1,6 @@ --extra-index-url https://download.pytorch.org/whl/cpu torch -torchvision +torchvision; platform_machine != "s390x" transformers huggingface-hub accelerate diff --git a/examples/model-conversion/scripts/causal/run-casual-gen-embeddings-org.py b/examples/model-conversion/scripts/causal/run-casual-gen-embeddings-org.py index b94bec4e765c..cb840dd5504b 100755 --- a/examples/model-conversion/scripts/causal/run-casual-gen-embeddings-org.py +++ b/examples/model-conversion/scripts/causal/run-casual-gen-embeddings-org.py @@ -2,12 +2,15 @@ import argparse import os +import sys import importlib import torch import numpy as np from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM -from pathlib import Path + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) +from utils.common import save_output_data unreleased_model_name = os.getenv('UNRELEASED_MODEL_NAME') @@ -54,6 +57,7 @@ prompt = "Hello world today" input_ids = tokenizer(prompt, return_tensors="pt").input_ids # ty: ignore[call-non-callable] +token_ids = input_ids[0].cpu().tolist() print(f"Input tokens: {input_ids}") print(f"Input text: {repr(prompt)}") print(f"Tokenized: {tokenizer.convert_ids_to_tokens(input_ids[0])}") # ty: ignore[unresolved-attribute] @@ -74,21 +78,8 @@ print(f"Hidden dimension: {token_embeddings.shape[-1]}") print(f"Number of tokens: {token_embeddings.shape[0]}") - # Save raw token embeddings - data_dir = Path("data") - data_dir.mkdir(exist_ok=True) - bin_filename = data_dir / f"pytorch-{model_name}-embeddings.bin" - txt_filename = data_dir / f"pytorch-{model_name}-embeddings.txt" - - # Save all token embeddings as binary print(token_embeddings) - token_embeddings.astype(np.float32).tofile(bin_filename) - - # Save as text for inspection - with open(txt_filename, "w") as f: - for i, embedding in enumerate(token_embeddings): - for j, val in enumerate(embedding): - f.write(f"{i} {j} {val:.6f}\n") + save_output_data(token_embeddings, token_ids, prompt, model_name, type_suffix="-embeddings") # Print embeddings per token in the requested format print("\nToken embeddings:") @@ -110,5 +101,3 @@ for i, token in enumerate(tokens): print(f" Token {i}: {repr(token)}") - print(f"Saved bin logits to: {bin_filename}") - print(f"Saved txt logist to: {txt_filename}") diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 2924fdbe9884..cc3f8cd36e35 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -154,6 +154,8 @@ extern "C" { bool buffer_from_host_ptr; // event synchronization bool events; + // mmap is supported for loading + bool mmap_support; }; // all the device properties diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index a5a3a58ad054..7654ea1f30f6 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -132,6 +132,7 @@ static void ggml_backend_meta_device_get_props(ggml_backend_dev_t dev, ggml_back /* .host_buffer = */ false, // Not implemented. /* .buffer_from_host_ptr = */ false, // Not implemented. /* .events = */ false, // Not implemented. + /* .mmap_support = */ true, }; for (ggml_backend_dev_t simple_dev : meta_dev_ctx->simple_devs) { ggml_backend_dev_props tmp_props; @@ -140,6 +141,7 @@ static void ggml_backend_meta_device_get_props(ggml_backend_dev_t dev, ggml_back props->caps.host_buffer = props->caps.host_buffer && tmp_props.caps.host_buffer; props->caps.buffer_from_host_ptr = props->caps.buffer_from_host_ptr && tmp_props.caps.buffer_from_host_ptr; props->caps.events = props->caps.events && tmp_props.caps.events; + props->caps.mmap_support = props->caps.mmap_support && tmp_props.caps.mmap_support; } } diff --git a/ggml/src/ggml-blas/ggml-blas.cpp b/ggml/src/ggml-blas/ggml-blas.cpp index 9745fa29f5db..e4b5bd254747 100644 --- a/ggml/src/ggml-blas/ggml-blas.cpp +++ b/ggml/src/ggml-blas/ggml-blas.cpp @@ -367,6 +367,7 @@ static void ggml_backend_blas_device_get_props(ggml_backend_dev_t dev, struct gg /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ true, /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-cann/ggml-cann.cpp b/ggml/src/ggml-cann/ggml-cann.cpp index 5f51ea3bb3c8..ffa361af4e8b 100644 --- a/ggml/src/ggml-cann/ggml-cann.cpp +++ b/ggml/src/ggml-cann/ggml-cann.cpp @@ -2815,6 +2815,7 @@ static void ggml_backend_cann_device_get_props(ggml_backend_dev_t dev, ggml_back /* .host_buffer = */ host_buffer, /* .buffer_from_host_ptr = */ false, /* .events = */ true, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 16cc5116c545..c0c9aa3cf09c 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -397,6 +397,7 @@ static void ggml_backend_cpu_device_get_props(ggml_backend_dev_t dev, struct ggm /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ true, /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 05e8d7f73684..1d4f4dfbd214 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -4801,6 +4801,7 @@ static void ggml_backend_cuda_device_get_props(ggml_backend_dev_t dev, ggml_back /* .host_buffer = */ host_buffer, /* .buffer_from_host_ptr = */ false, /* .events = */ events, + /* .mmap_support = */ props->type != GGML_BACKEND_DEVICE_TYPE_IGPU, }; } diff --git a/ggml/src/ggml-et/ggml-et.cpp b/ggml/src/ggml-et/ggml-et.cpp index b30209095672..e8482f734621 100644 --- a/ggml/src/ggml-et/ggml-et.cpp +++ b/ggml/src/ggml-et/ggml-et.cpp @@ -1646,6 +1646,7 @@ static void ggml_backend_et_device_get_props(ggml_backend_dev_t dev, struct ggml /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ false, /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index bdb8af0820a3..f80c60a500bd 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -3930,6 +3930,7 @@ static void ggml_backend_hexagon_device_get_props(ggml_backend_dev_t dev, struct /* .host_buffer = */ (bool) opt_hostbuf, /* .buffer_from_host_ptr = */ false, /* .events = */ false, + /* .mmap_support = */ false, }; } diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp index a1003b3acff8..ef3c92f27127 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp @@ -681,6 +681,7 @@ static void ggml_backend_metal_device_get_props(ggml_backend_dev_t dev, ggml_bac /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ true, /* .events = */ true, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index fc0fce0d780a..19dca4ced0b8 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -73,6 +73,7 @@ typedef const void * (*get_adreno_bin_kernel_func_t)( //------------------------------------------------------------------------------ bool ggml_cl_compute_forward(ggml_backend_t backend, struct ggml_tensor * tensor); + static bool ggml_cl_is_q4_0_soa(const ggml_tensor * tensor); static bool ggml_cl_is_q8_0_soa(const ggml_tensor * tensor); static void ggml_cl_mul_mat(ggml_backend_t backend, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); @@ -4629,6 +4630,23 @@ static std::string ggml_opencl_fa_compile_opts(ggml_backend_opencl_context * bac if (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E) { opts += " -D FA_C8_NO_SG_PIN"; } + // Transposed K tile in local memory: the KV rows the QK loop walks together become + // adjacent, so a group of them is ONE 128-bit local read instead of several narrow + // ones. The QK loop is LDS-read-issue-bound (a wrong-math probe that kept every FMA/dp4a + // but removed the LDS reads ran the kernel ~40% faster), so this is worth up to +26% on + // fa=1 prefill. Output is bit-identical -- only the layout moves. + // + // DK <= 128 only. At DK=256 (gemma-3-4b) it measures 1-2% NEGATIVE and reproduces across + // rounds; padding the row stride does not recover it, so the cause is not a simple bank + // conflict and the wider tile does not want this layout. + // + // Default on within that gate; GGML_OPENCL_FA_K_LDS_T=0 restores the row-major tile. + { + const char * e = getenv("GGML_OPENCL_FA_K_LDS_T"); + if ((e == nullptr || e[0] != '0') && cfg->dk <= 128) { + opts += " -D FA_K_LDS_T"; + } + } return opts; } @@ -10777,6 +10795,7 @@ static void ggml_backend_opencl_device_get_props(ggml_backend_dev_t dev, struct /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ false, /* .events = */ false, + /* .mmap_support = */ false, }; } diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl index 6e43ee81e73b..bf7695a2c1d7 100644 --- a/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl +++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl @@ -211,7 +211,30 @@ __kernel void FA_TILE_NAME( float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1); +#ifdef FA_K_LDS_T + // K tile transposed: [dk vec][kv row] instead of [kv row][dk vec]. + // + // The QK loop walks 2 or 4 KV rows at a time against the same dk element. Row-major + // those are DK_VEC half4s apart, so each is its own 64-bit local read. Transposed they + // are adjacent, so a pair is one 128-bit read -- half the LDS issues for the same bytes, + // no extra registers, arithmetic untouched. + // + // This kernel looked like it should be FMA-bound (a half4 mad does ~4 ALU ops per LDS + // read, unlike the 1:1 of the dp4a loop), but it is NOT: a wrong-math probe that kept + // every FMA and removed the LDS reads ran it 38.6% faster (18.92 -> 11.62 ms/op). + // Explicitly 16-byte aligned: FA_LK_PAIR below reads two adjacent half4 as one float4, + // and the element type only obliges the compiler to align this array to 8. The indices + // are even so the offset is a multiple of 16, but the base has to be too, and relying + // on the compiler to over-align it is relying on luck. + __local KV_DATA_TYPE4 l_k[DK_VEC][BLOCK_N] __attribute__((aligned(16))); +#define FA_LK(ROW, C) l_k[C][ROW] + // Two adjacent KV rows as one 128-bit local read (half4 pair == 16 B). j is even and + // BLOCK_N is even, so &l_k[c][j] is 16 B past a 16 B-aligned base. +#define FA_LK_PAIR(C, J) as_half8(*(__local const float4 *)(&l_k[C][J])) +#else __local KV_DATA_TYPE4 l_k[BLOCK_N][DK_VEC]; +#define FA_LK(ROW, C) l_k[ROW][C] +#endif __local KV_DATA_TYPE4 l_v[BLOCK_N][DV_VEC]; #if N_SPLIT > 1 && !defined(HAS_SUBGROUP_SHUFFLE) @@ -254,17 +277,17 @@ __kernel void FA_TILE_NAME( #ifdef FA_K_IMG if (use_kv_pad) { const ulong k_row_offset = batch_idx * k_tile_nb3 + head_kv_idx * k_tile_nb2 + k_row_idx * k_nb1; - l_k[row][col] = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col]; + FA_LK(row, col) = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col]; } else { const int k_row_px = batch_idx * k_pitch_px_batch + head_kv_idx * k_pitch_px_head + k_row_idx * k_pitch_px_row; - l_k[row][col] = read_imageh(k_img, k_row_px + col); + FA_LK(row, col) = read_imageh(k_img, k_row_px + col); } #else const ulong k_row_offset = batch_idx * k_tile_nb3 + head_kv_idx * k_tile_nb2 + k_row_idx * k_nb1; - l_k[row][col] = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col]; + FA_LK(row, col) = ((__global KV_DATA_TYPE4*)(k_tile_base + k_row_offset))[col]; #endif } else { - l_k[row][col] = (KV_DATA_TYPE4)(0.0h); + FA_LK(row, col) = (KV_DATA_TYPE4)(0.0h); } } for (int i = tid; i < BLOCK_N * DV_VEC; i += WG_SIZE) { @@ -292,8 +315,15 @@ __kernel void FA_TILE_NAME( FA_UNROLL for (int k = 0; k < SPLIT_DK_VEC; k++) { const ACC_TYPE4 qk = q_priv[k]; +#if defined(FA_K_LDS_T) + // 2 KV rows adjacent in the transposed tile: one 128-bit local read. + const half8 kk = FA_LK_PAIR(dk_off + k, j); + ACC_TYPE4 dot0 = qk * CONVERT_KV_ACC4(kk.lo); + ACC_TYPE4 dot1 = qk * CONVERT_KV_ACC4(kk.hi); +#else ACC_TYPE4 dot0 = qk * CONVERT_KV_ACC4(l_k[j ][dk_off + k]); ACC_TYPE4 dot1 = qk * CONVERT_KV_ACC4(l_k[j+1][dk_off + k]); +#endif partial0 += dot0.s0 + dot0.s1 + dot0.s2 + dot0.s3; partial1 += dot1.s0 + dot1.s1 + dot1.s2 + dot1.s3; } @@ -359,7 +389,7 @@ __kernel void FA_TILE_NAME( ACC_TYPE4 dot_acc = (ACC_TYPE4)(0.0f); FA_UNROLL for (int k = 0; k < SPLIT_DK_VEC; k++) { - dot_acc = mad(q_priv[k], CONVERT_KV_ACC4(l_k[j][dk_off + k]), dot_acc); + dot_acc = mad(q_priv[k], CONVERT_KV_ACC4(FA_LK(j, dk_off + k)), dot_acc); } local_partial[j][tid] = dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3; @@ -452,10 +482,21 @@ __kernel void FA_TILE_NAME( FA_UNROLL for (int k = 0; k < DK_VEC; k++) { const ACC_TYPE4 qk = q_priv[k]; +#if defined(FA_K_LDS_T) + // 4 KV rows adjacent in the transposed tile: two 128-bit local reads + // instead of four 64-bit ones. + const half8 kk01 = FA_LK_PAIR(k, j); + const half8 kk23 = FA_LK_PAIR(k, j + 2); + dot_acc0 = mad(qk, CONVERT_KV_ACC4(kk01.lo), dot_acc0); + dot_acc1 = mad(qk, CONVERT_KV_ACC4(kk01.hi), dot_acc1); + dot_acc2 = mad(qk, CONVERT_KV_ACC4(kk23.lo), dot_acc2); + dot_acc3 = mad(qk, CONVERT_KV_ACC4(kk23.hi), dot_acc3); +#else dot_acc0 = mad(qk, CONVERT_KV_ACC4(l_k[j][k]), dot_acc0); dot_acc1 = mad(qk, CONVERT_KV_ACC4(l_k[j+1][k]), dot_acc1); dot_acc2 = mad(qk, CONVERT_KV_ACC4(l_k[j+2][k]), dot_acc2); dot_acc3 = mad(qk, CONVERT_KV_ACC4(l_k[j+3][k]), dot_acc3); +#endif } ACC_TYPE s0 = (dot_acc0.s0 + dot_acc0.s1 + dot_acc0.s2 + dot_acc0.s3) * scale; ACC_TYPE s1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale; diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl index 95d215971e00..48adba4f725b 100644 --- a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl +++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl @@ -1631,8 +1631,25 @@ __kernel void flash_attn_f32_q4_0( float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1); #ifdef FA_HAVE_INT_DOT +// Accessors so the staging code is layout-agnostic. +#ifdef FA_K_LDS_T +#define FA_K_PACKED(ROW, IDX) l_k_packed[IDX][ROW] +#define FA_K_SCALE(ROW, BLK) l_k_scale[BLK][ROW] +#else +#define FA_K_PACKED(ROW, IDX) l_k_packed[ROW][IDX] +#define FA_K_SCALE(ROW, BLK) l_k_scale[ROW][BLK] +#endif + +#ifdef FA_K_LDS_T + // K tile transposed: the 4 KV rows the QK loop walks together become adjacent, so each + // (block, group) step is ONE 128-bit local read instead of four 32-bit ones. The QK + // loop is LDS-read-issue-bound. + __local uint l_k_packed[DK_Q4_BLOCKS_PREFILL * 8][BLOCK_N]; + __local float l_k_scale [DK_Q4_BLOCKS_PREFILL][BLOCK_N]; +#else __local uint l_k_packed[BLOCK_N][DK_Q4_BLOCKS_PREFILL * 8]; __local float l_k_scale [BLOCK_N][DK_Q4_BLOCKS_PREFILL]; +#endif #else __local half4 l_k[BLOCK_N][DK_VEC]; #endif @@ -1660,17 +1677,17 @@ __kernel void flash_attn_f32_q4_0( const global char * blk_ptr = k_base + k_row_off + blk * Q4_0_BLOCK_SIZE; const float df = (float) vload_half(0, (const global half *) blk_ptr); const global uchar * qs = (const global uchar *)(blk_ptr + 2); - l_k_scale[row][blk] = df; + FA_K_SCALE(row, blk) = df; uint k_packed[8]; pack_q4_0_nibbles(qs, k_packed); #pragma unroll for (int j = 0; j < 8; ++j) { - l_k_packed[row][blk * 8 + j] = k_packed[j]; + FA_K_PACKED(row, blk * 8 + j) = k_packed[j]; } } else { - l_k_scale[row][blk] = 0.0f; + FA_K_SCALE(row, blk) = 0.0f; #pragma unroll - for (int j = 0; j < 8; ++j) l_k_packed[row][blk * 8 + j] = 0u; + for (int j = 0; j < 8; ++j) FA_K_PACKED(row, blk * 8 + j) = 0u; } } #else @@ -1760,6 +1777,19 @@ __kernel void flash_attn_f32_q4_0( for (int b_local = 0; b_local < SPLIT_DK_Q4_BLOCKS; ++b_local) { const int b = k_blk_base + b_local; int sum0 = 0, sum1 = 0, sum2 = 0, sum3 = 0; +#ifdef FA_K_LDS_T + // 4 KV rows are adjacent in the transposed tile: one 128-bit local + // read per (block, group) instead of four 32-bit ones. + #pragma unroll + for (int g = 0; g < 8; ++g) { + const uint qp = q_packed_pf[b_local * 8 + g]; + const uint4 kq4 = vload4(0, &l_k_packed[b * 8 + g][j]); + sum0 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s0, sum0); + sum1 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s1, sum1); + sum2 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s2, sum2); + sum3 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s3, sum3); + } +#else #pragma unroll for (int g = 0; g < 8; ++g) { const uint qp = q_packed_pf[b_local * 8 + g]; @@ -1768,12 +1798,21 @@ __kernel void flash_attn_f32_q4_0( sum2 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+2][b * 8 + g], sum2); sum3 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+3][b * 8 + g], sum3); } +#endif const float qd = q_d_pf[b_local]; const int q_sum = q_sum_pf[b_local]; +#ifdef FA_K_LDS_T + const float4 ks4 = vload4(0, &l_k_scale[b][j]); + s0 += (float)(sum0 - 8 * q_sum) * qd * ks4.s0; + s1 += (float)(sum1 - 8 * q_sum) * qd * ks4.s1; + s2 += (float)(sum2 - 8 * q_sum) * qd * ks4.s2; + s3 += (float)(sum3 - 8 * q_sum) * qd * ks4.s3; +#else s0 += (float)(sum0 - 8 * q_sum) * qd * l_k_scale[j ][b]; s1 += (float)(sum1 - 8 * q_sum) * qd * l_k_scale[j+1][b]; s2 += (float)(sum2 - 8 * q_sum) * qd * l_k_scale[j+2][b]; s3 += (float)(sum3 - 8 * q_sum) * qd * l_k_scale[j+3][b]; +#endif } #else ACC_TYPE4 dot_acc0 = (ACC_TYPE4)(0.0f); diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl index 7e89ed0bd8f1..f50912d21101 100644 --- a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl +++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl @@ -1393,8 +1393,31 @@ __kernel void flash_attn_f32_q8_0( float slope = get_alibi_slope(max_bias, head_idx, n_head_log2, m0, m1); #ifdef FA_HAVE_INT_DOT +// Accessors so the staging code is layout-agnostic. +#ifdef FA_K_LDS_T +#define FA_K_PACKED(ROW, IDX) l_k_packed[IDX][ROW] +#define FA_K_SCALE(ROW, BLK) l_k_scale[BLK][ROW] +#else +#define FA_K_PACKED(ROW, IDX) l_k_packed[ROW][IDX] +#define FA_K_SCALE(ROW, BLK) l_k_scale[ROW][BLK] +#endif + +#ifdef FA_K_LDS_T + // K tile transposed: [block*8 + g][kv row] instead of [kv row][block*8 + g]. + // + // The QK loop walks 4 KV rows at a time against the same (b, g), so in the original + // layout those 4 values are BLOCK_N*8 uints apart and cost 4 separate 32-bit local + // reads. Transposed they are adjacent, so they are one 128-bit read -- 4x fewer LDS + // issues for the same bytes and no extra registers. That matters because the QK loop + // is LDS-read-issue-bound: a wrong-math probe that kept every dp4a but cut the LDS + // reads ran the whole kernel 41% faster (18.51 -> 10.91 ms/op), and deleting QK + // outright only reached 10.88 -- i.e. essentially ALL of QK's cost is these reads. + __local uint l_k_packed[DK_Q8_BLOCKS_PREFILL * 8][BLOCK_N]; + __local float l_k_scale [DK_Q8_BLOCKS_PREFILL][BLOCK_N]; +#else __local uint l_k_packed[BLOCK_N][DK_Q8_BLOCKS_PREFILL * 8]; __local float l_k_scale [BLOCK_N][DK_Q8_BLOCKS_PREFILL]; +#endif #else __local half4 l_k[BLOCK_N][DK_VEC]; #endif @@ -1427,7 +1450,7 @@ __kernel void flash_attn_f32_q8_0( const global char * blk_ptr = k_base + k_row_off + blk * Q8_0_BLOCK_SIZE; const float df = (float) vload_half(0, (const global half *) blk_ptr); const global uchar * qs = (const global uchar *)(blk_ptr + 2); - l_k_scale[row][blk] = df; + FA_K_SCALE(row, blk) = df; #pragma unroll for (int j = 0; j < 8; ++j) { uint k_packed = @@ -1435,12 +1458,12 @@ __kernel void flash_attn_f32_q8_0( ((uint) qs[j*4 + 1]) << 8 | ((uint) qs[j*4 + 2]) << 16 | ((uint) qs[j*4 + 3]) << 24; - l_k_packed[row][blk * 8 + j] = k_packed; + FA_K_PACKED(row, blk * 8 + j) = k_packed; } } else { - l_k_scale[row][blk] = 0.0f; + FA_K_SCALE(row, blk) = 0.0f; #pragma unroll - for (int j = 0; j < 8; ++j) l_k_packed[row][blk * 8 + j] = 0u; + for (int j = 0; j < 8; ++j) FA_K_PACKED(row, blk * 8 + j) = 0u; } } #else @@ -1556,6 +1579,19 @@ __kernel void flash_attn_f32_q8_0( for (int b_local = 0; b_local < SPLIT_DK_Q8_BLOCKS; ++b_local) { const int b = k_blk_base + b_local; int sum0 = 0, sum1 = 0, sum2 = 0, sum3 = 0; +#if defined(FA_K_LDS_T) + // The 4 KV rows are adjacent in the transposed tile, so each (b, g) + // step is ONE 128-bit local read instead of four 32-bit ones. + #pragma unroll + for (int g = 0; g < 8; ++g) { + const uint qp = q_packed_pf[b_local * 8 + g]; + const uint4 kq4 = vload4(0, &l_k_packed[b * 8 + g][j]); + sum0 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s0, sum0); + sum1 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s1, sum1); + sum2 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s2, sum2); + sum3 = dot_acc_sat_4x8packed_ss_int(qp, kq4.s3, sum3); + } +#else #pragma unroll for (int g = 0; g < 8; ++g) { const uint qp = q_packed_pf[b_local * 8 + g]; @@ -1564,11 +1600,20 @@ __kernel void flash_attn_f32_q8_0( sum2 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+2][b * 8 + g], sum2); sum3 = dot_acc_sat_4x8packed_ss_int(qp, l_k_packed[j+3][b * 8 + g], sum3); } +#endif const float qd = q_d_pf[b_local]; +#ifdef FA_K_LDS_T + const float4 ks4 = vload4(0, &l_k_scale[b][j]); + s0 += (float)sum0 * qd * ks4.s0; + s1 += (float)sum1 * qd * ks4.s1; + s2 += (float)sum2 * qd * ks4.s2; + s3 += (float)sum3 * qd * ks4.s3; +#else s0 += (float)sum0 * qd * l_k_scale[j ][b]; s1 += (float)sum1 * qd * l_k_scale[j+1][b]; s2 += (float)sum2 * qd * l_k_scale[j+2][b]; s3 += (float)sum3 * qd * l_k_scale[j+3][b]; +#endif } #else ACC_TYPE4 dot_acc0 = (ACC_TYPE4)(0.0f); diff --git a/ggml/src/ggml-openvino/ggml-openvino.cpp b/ggml/src/ggml-openvino/ggml-openvino.cpp index 0e7501fefe38..dfe80e6c86e7 100644 --- a/ggml/src/ggml-openvino/ggml-openvino.cpp +++ b/ggml/src/ggml-openvino/ggml-openvino.cpp @@ -763,6 +763,7 @@ static void ggml_backend_openvino_device_get_props(ggml_backend_dev_t dev, ggml_ /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ false, /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index 17c53a5f049e..e9de0d0aa98a 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -1881,6 +1881,7 @@ static void ggml_backend_rpc_device_get_props(ggml_backend_dev_t dev, struct ggm /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ false, /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 18d58782ebff..8a19f648bf07 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -5649,6 +5649,7 @@ static void ggml_backend_sycl_device_get_props(ggml_backend_dev_t dev, ggml_back /* .host_buffer = */ host_buffer, /* .buffer_from_host_ptr = */ false, /* .events = */ events, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-virtgpu/backend/backend-dispatched-device.cpp b/ggml/src/ggml-virtgpu/backend/backend-dispatched-device.cpp index c7acb8b51ce7..87872df1c7bc 100644 --- a/ggml/src/ggml-virtgpu/backend/backend-dispatched-device.cpp +++ b/ggml/src/ggml-virtgpu/backend/backend-dispatched-device.cpp @@ -111,6 +111,7 @@ uint32_t backend_device_get_props(apir_encoder * enc, apir_decoder * dec, virgl_ apir_encode_bool_t(enc, &props.caps.host_buffer); apir_encode_bool_t(enc, &props.caps.buffer_from_host_ptr); apir_encode_bool_t(enc, &props.caps.events); + apir_encode_bool_t(enc, &props.caps.mmap_support); return 0; } diff --git a/ggml/src/ggml-virtgpu/backend/shared/api_remoting.h b/ggml/src/ggml-virtgpu/backend/shared/api_remoting.h index 6bf97e8a3a24..a5ef3ea476d0 100644 --- a/ggml/src/ggml-virtgpu/backend/shared/api_remoting.h +++ b/ggml/src/ggml-virtgpu/backend/shared/api_remoting.h @@ -7,7 +7,7 @@ #include #define APIR_PROTOCOL_MAJOR 0 -#define APIR_PROTOCOL_MINOR 1 +#define APIR_PROTOCOL_MINOR 2 #define APIR_HANDSHAKE_MAGIC 0xab1e diff --git a/ggml/src/ggml-virtgpu/ggml-backend-buffer-type.cpp b/ggml/src/ggml-virtgpu/ggml-backend-buffer-type.cpp index 8fa20ff43bd5..d5bdc993b464 100644 --- a/ggml/src/ggml-virtgpu/ggml-backend-buffer-type.cpp +++ b/ggml/src/ggml-virtgpu/ggml-backend-buffer-type.cpp @@ -11,9 +11,9 @@ static ggml_backend_buffer_t ggml_backend_remoting_buffer_type_alloc_buffer(ggml context->gpu = gpu; - bool async__unused, host_buffer__unused, events__unused; + bool async__unused, host_buffer__unused, events__unused, mmap_support__unused; bool buffer_from_host_ptr; - apir_device_get_props(gpu, &async__unused, &host_buffer__unused, &buffer_from_host_ptr, &events__unused); + apir_device_get_props(gpu, &async__unused, &host_buffer__unused, &buffer_from_host_ptr, &events__unused, &mmap_support__unused); if (buffer_from_host_ptr) { context->apir_context = apir_device_buffer_from_ptr(gpu, size, size); diff --git a/ggml/src/ggml-virtgpu/ggml-backend-device.cpp b/ggml/src/ggml-virtgpu/ggml-backend-device.cpp index a978812cd908..987ce9dd110c 100644 --- a/ggml/src/ggml-virtgpu/ggml-backend-device.cpp +++ b/ggml/src/ggml-virtgpu/ggml-backend-device.cpp @@ -65,7 +65,7 @@ static void ggml_backend_remoting_device_get_props(ggml_backend_dev_t dev, ggml_ virtgpu * gpu = DEV_TO_GPU(dev); apir_device_get_props(gpu, &props->caps.async, &props->caps.host_buffer, &props->caps.buffer_from_host_ptr, - &props->caps.events); + &props->caps.events, &props->caps.mmap_support); props->caps.buffer_from_host_ptr = false; props->caps.async = false; diff --git a/ggml/src/ggml-virtgpu/virtgpu-forward-device.cpp b/ggml/src/ggml-virtgpu/virtgpu-forward-device.cpp index 9f513c138dd2..864264f213b5 100644 --- a/ggml/src/ggml-virtgpu/virtgpu-forward-device.cpp +++ b/ggml/src/ggml-virtgpu/virtgpu-forward-device.cpp @@ -144,7 +144,8 @@ void apir_device_get_props(virtgpu * gpu, bool * async, bool * host_buffer, bool * buffer_from_host_ptr, - bool * events) { + bool * events, + bool * mmap_support) { apir_encoder * encoder; apir_decoder * decoder; ApirForwardReturnCode ret; @@ -157,6 +158,7 @@ void apir_device_get_props(virtgpu * gpu, apir_decode_bool_t(decoder, host_buffer); apir_decode_bool_t(decoder, buffer_from_host_ptr); apir_decode_bool_t(decoder, events); + apir_decode_bool_t(decoder, mmap_support); remote_call_finish(gpu, encoder, decoder); diff --git a/ggml/src/ggml-virtgpu/virtgpu-forward.gen.h b/ggml/src/ggml-virtgpu/virtgpu-forward.gen.h index 44b0ad1ffa1d..da28aa5f9047 100644 --- a/ggml/src/ggml-virtgpu/virtgpu-forward.gen.h +++ b/ggml/src/ggml-virtgpu/virtgpu-forward.gen.h @@ -13,7 +13,8 @@ void apir_device_get_props(struct virtgpu * gpu, bool * async, bool * host_buffer, bool * buffer_from_host_ptr, - bool * events); + bool * events, + bool * mmap_support); apir_buffer_context_t apir_device_buffer_from_ptr(struct virtgpu * gpu, size_t size, size_t max_tensor_size); /* buffer-type */ diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index a923755f9edd..45fa97f81297 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -17891,6 +17891,7 @@ static void ggml_backend_vk_device_get_props(ggml_backend_dev_t dev, struct ggml /* .host_buffer = */ true, /* .buffer_from_host_ptr = */ false, /* .events = */ true, + /* .mmap_support = */ !ctx->is_integrated_gpu, }; } diff --git a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp index 35a55ecaf644..0604e1c2b87b 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp @@ -2815,11 +2815,25 @@ class ggml_webgpu_shader_lib { key.common.v_direct &= decisions.use_sg_matrix && key.common.v_type == GGML_TYPE_F16; key.use_sg_matrix = decisions.use_sg_matrix; - const uint32_t max_kv_tile = ggml_webgpu_flash_attn_max_kv_tile( + uint32_t max_kv_tile = ggml_webgpu_flash_attn_max_kv_tile( context.wg_mem_limit_bytes, decisions.q_tile, decisions.use_sg_matrix ? context.sg_mat_n : 1u, key.common.head_dim_qk, key.common.head_dim_v, key.common.has_mask, key.common.k_direct || key.common.v_direct); - GGML_ASSERT(max_kv_tile > 0); + + // WorkGroup storage size isn't enough for some params with subgroup matrices path (ref. https://github.com/ggml-org/llama.cpp/pull/26566) + if (max_kv_tile == 0) { + GGML_ASSERT(decisions.use_sg_matrix); + // switch to flash_attn_reg_tile path + decisions.use_sg_matrix = false; + decisions.q_tile = GGML_WEBGPU_FLASH_ATTN_TILE_Q_TILE; + key.common.k_direct = false; + key.common.v_direct = false; + key.use_sg_matrix = false; + max_kv_tile = ggml_webgpu_flash_attn_max_kv_tile( + context.wg_mem_limit_bytes, decisions.q_tile, 1u, key.common.head_dim_qk, key.common.head_dim_v, + key.common.has_mask, key.common.k_direct || key.common.v_direct); + GGML_ASSERT(max_kv_tile > 0); + } decisions.kv_tile = decisions.use_sg_matrix ? std::min(max_kv_tile, context.sg_mat_n * GGML_WEBGPU_FLASH_ATTN_PREFERRED_KV_SG_TILES) : @@ -2993,6 +3007,10 @@ class ggml_webgpu_shader_lib { defines.push_back("SRC_F16"); variant += "_f16"; break; + case GGML_TYPE_I32: + defines.push_back("SRC_I32"); + variant += "_i32"; + break; default: GGML_ABORT("Unsupported src type for cpy shader"); } diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index ba4b91695fae..6741752b3611 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -3942,6 +3942,7 @@ static void ggml_backend_webgpu_device_get_props(ggml_backend_dev_t dev, struct /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ false, /* .events = */ false, + /* .mmap_support = */ true, }; } @@ -4283,9 +4284,8 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const break; case GGML_OP_CPY: case GGML_OP_CONT: - supports_op = ((op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) && - (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16)) || - (op->type == GGML_TYPE_I32 && src0->type == GGML_TYPE_F32); + supports_op = (op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_I32) && + (src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_I32); break; case GGML_OP_SET: supports_op = src0->type == src1->type && src0->type == op->type && diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/cpy.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/cpy.wgsl index 67f1dc0928f8..0d0d81ab650f 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/cpy.wgsl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/cpy.wgsl @@ -4,6 +4,8 @@ enable f16; #define SRC_TYPE f32 #elif defined(SRC_F16) #define SRC_TYPE f16 +#elif defined(SRC_I32) +#define SRC_TYPE i32 #endif #ifdef DST_F32 diff --git a/ggml/src/ggml-zdnn/ggml-zdnn.cpp b/ggml/src/ggml-zdnn/ggml-zdnn.cpp index 639b818d128e..4007ac9dfc7d 100644 --- a/ggml/src/ggml-zdnn/ggml-zdnn.cpp +++ b/ggml/src/ggml-zdnn/ggml-zdnn.cpp @@ -487,7 +487,8 @@ static void ggml_backend_zdnn_device_get_props(ggml_backend_dev_t dev, ggml_back /* .async = */ false, /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ false, - /* .events = */ false + /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/ggml/src/ggml-zendnn/ggml-zendnn.cpp b/ggml/src/ggml-zendnn/ggml-zendnn.cpp index e6a9b51b7925..ec7ce233145a 100644 --- a/ggml/src/ggml-zendnn/ggml-zendnn.cpp +++ b/ggml/src/ggml-zendnn/ggml-zendnn.cpp @@ -654,7 +654,8 @@ static void ggml_backend_zendnn_device_get_props(ggml_backend_dev_t dev, struct /* .async = */ false, /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ true, - /* .events = */ false + /* .events = */ false, + /* .mmap_support = */ true, }; } diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 8d0cad59b2cd..98c4fa169143 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -407,6 +407,8 @@ class Projector: class ClipGenAudio: PROJECTOR_TYPE = "clip.gen.audio.projector_type" # for mixed modality models + # name of the weight variant, for settings that are not in the checkpoint + MODEL_VARIANT = "clip.gen.audio.model_variant" EMBEDDING_LENGTH = "clip.gen.audio.embedding_length" FEED_FORWARD_LENGTH = "clip.gen.audio.feed_forward_length" BLOCK_COUNT = "clip.gen.audio.block_count" @@ -581,6 +583,7 @@ class MODEL_ARCH(IntEnum): MELLUM = auto() NANBEIGE = auto() QWEN3TTS = auto() + POCKETTTS = auto() class VISION_PROJECTOR_TYPE(IntEnum): @@ -1040,6 +1043,38 @@ class MODEL_TENSOR(IntEnum): A_GEN_WAV_DAC_RES_CONV2 = auto() # DAC residual unit, pointwise causal conv A_GEN_WAV_DAC_POST_SNAKE = auto() # DAC final SnakeBeta A_GEN_WAV_DAC_POST_CONV = auto() # DAC conv_post -> 1-channel PCM + # pocket-tts: SEANet encoder (speaker path) and decoder (a.gen.wav path) + A_ENC_SEANET_CONV_IN = auto() + A_ENC_SEANET_CONV_OUT = auto() + A_ENC_SEANET_RES_CONV1 = auto() # residual unit, dilated conv + A_ENC_SEANET_RES_CONV2 = auto() # residual unit, pointwise conv + A_ENC_SEANET_SCALE_CONV = auto() # strided downsample conv + A_ENC_ATTN_SCALE = auto() # layer scale (gamma) on the attn output + A_ENC_FFN_SCALE_LS = auto() # layer scale (gamma) on the FFN output + A_ENC_SPEAKER_PROJ = auto() # voice latent -> backbone embd + A_GEN_FLOW_INPUT_PROJ = auto() + A_GEN_FLOW_COND_EMBD = auto() + A_GEN_FLOW_TIME_FREQS = auto() # timestep embedder, stored cos/sin frequencies + A_GEN_FLOW_TIME_UP = auto() + A_GEN_FLOW_TIME_DOWN = auto() + A_GEN_FLOW_TIME_NORM = auto() # RMSNorm alpha + A_GEN_FLOW_BLK_NORM = auto() # AdaLN res block, in_ln + A_GEN_FLOW_BLK_UP = auto() + A_GEN_FLOW_BLK_DOWN = auto() + A_GEN_FLOW_BLK_ADA = auto() # AdaLN modulation, -> shift/scale/gate + A_GEN_FLOW_FINAL_ADA = auto() # final layer AdaLN modulation, -> shift/scale + A_GEN_FLOW_FINAL_PROJ = auto() + A_GEN_OUT_EOS = auto() # end-of-speech head on the backbone hidden state + A_GEN_INPUT_LINEAR = auto() # generated latent -> backbone embd + A_GEN_EMB_MEAN = auto() # latent denormalization stats + A_GEN_EMB_STD = auto() + A_GEN_WAV_QUANT_OUT = auto() # DummyQuantizer output_proj, latent -> decoder dim + A_GEN_WAV_UPSAMPLE = auto() # frame rate -> encoder frame rate, depthwise convtr + A_GEN_WAV_SEANET_CONV_IN = auto() + A_GEN_WAV_SEANET_CONV_OUT = auto() # -> 1-channel PCM + A_GEN_WAV_SEANET_RES_CONV1 = auto() + A_GEN_WAV_SEANET_RES_CONV2 = auto() + A_GEN_WAV_SEANET_SCALE_CONV = auto() # strided upsample convtr A_MMPROJ = auto() A_MMPROJ_FC = auto() A_MM_NORM_PRE = auto() @@ -1255,6 +1290,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.MELLUM: "mellum", MODEL_ARCH.NANBEIGE: "nanbeige", MODEL_ARCH.QWEN3TTS: "qwen3tts", + MODEL_ARCH.POCKETTTS: "pockettts", } VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = { @@ -1709,6 +1745,37 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV2: "a.gen.wav.dac.blk.{bid}.res.{xid}.conv2", MODEL_TENSOR.A_GEN_WAV_DAC_POST_SNAKE: "a.gen.wav.dac.post_snake", MODEL_TENSOR.A_GEN_WAV_DAC_POST_CONV: "a.gen.wav.dac.post_conv", + MODEL_TENSOR.A_ENC_SEANET_CONV_IN: "a.seanet.conv_in", + MODEL_TENSOR.A_ENC_SEANET_CONV_OUT: "a.seanet.conv_out", + MODEL_TENSOR.A_ENC_SEANET_RES_CONV1: "a.seanet.blk.{bid}.res_conv1", + MODEL_TENSOR.A_ENC_SEANET_RES_CONV2: "a.seanet.blk.{bid}.res_conv2", + MODEL_TENSOR.A_ENC_SEANET_SCALE_CONV: "a.seanet.blk.{bid}.scale_conv", + MODEL_TENSOR.A_ENC_ATTN_SCALE: "a.blk.{bid}.ls1", + MODEL_TENSOR.A_ENC_FFN_SCALE_LS: "a.blk.{bid}.ls2", + MODEL_TENSOR.A_ENC_SPEAKER_PROJ: "a.speaker_proj", + MODEL_TENSOR.A_GEN_FLOW_INPUT_PROJ: "a.gen.flow.input_proj", + MODEL_TENSOR.A_GEN_FLOW_COND_EMBD: "a.gen.flow.cond_embd", + MODEL_TENSOR.A_GEN_FLOW_TIME_FREQS: "a.gen.flow.time.{bid}.freqs", + MODEL_TENSOR.A_GEN_FLOW_TIME_UP: "a.gen.flow.time.{bid}.up", + MODEL_TENSOR.A_GEN_FLOW_TIME_DOWN: "a.gen.flow.time.{bid}.down", + MODEL_TENSOR.A_GEN_FLOW_TIME_NORM: "a.gen.flow.time.{bid}.norm", + MODEL_TENSOR.A_GEN_FLOW_BLK_NORM: "a.gen.flow.blk.{bid}.norm", + MODEL_TENSOR.A_GEN_FLOW_BLK_UP: "a.gen.flow.blk.{bid}.up", + MODEL_TENSOR.A_GEN_FLOW_BLK_DOWN: "a.gen.flow.blk.{bid}.down", + MODEL_TENSOR.A_GEN_FLOW_BLK_ADA: "a.gen.flow.blk.{bid}.ada", + MODEL_TENSOR.A_GEN_FLOW_FINAL_ADA: "a.gen.flow.final.ada", + MODEL_TENSOR.A_GEN_FLOW_FINAL_PROJ: "a.gen.flow.final.proj", + MODEL_TENSOR.A_GEN_OUT_EOS: "a.gen.out_eos", + MODEL_TENSOR.A_GEN_INPUT_LINEAR: "a.gen.input_linear", + MODEL_TENSOR.A_GEN_EMB_MEAN: "a.gen.emb_mean", + MODEL_TENSOR.A_GEN_EMB_STD: "a.gen.emb_std", + MODEL_TENSOR.A_GEN_WAV_QUANT_OUT: "a.gen.wav.quant_out", + MODEL_TENSOR.A_GEN_WAV_UPSAMPLE: "a.gen.wav.upsample", + MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_IN: "a.gen.wav.seanet.conv_in", + MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_OUT: "a.gen.wav.seanet.conv_out", + MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV1: "a.gen.wav.seanet.blk.{bid}.res_conv1", + MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV2: "a.gen.wav.seanet.blk.{bid}.res_conv2", + MODEL_TENSOR.A_GEN_WAV_SEANET_SCALE_CONV: "a.gen.wav.seanet.blk.{bid}.scale_conv", MODEL_TENSOR.A_MMPROJ: "mm.a.mlp.{bid}", MODEL_TENSOR.A_MMPROJ_FC: "mm.a.fc", MODEL_TENSOR.A_MM_NORM_PRE: "mm.a.norm_pre", @@ -2020,6 +2087,37 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV2, MODEL_TENSOR.A_GEN_WAV_DAC_POST_SNAKE, MODEL_TENSOR.A_GEN_WAV_DAC_POST_CONV, + MODEL_TENSOR.A_ENC_SEANET_CONV_IN, + MODEL_TENSOR.A_ENC_SEANET_CONV_OUT, + MODEL_TENSOR.A_ENC_SEANET_RES_CONV1, + MODEL_TENSOR.A_ENC_SEANET_RES_CONV2, + MODEL_TENSOR.A_ENC_SEANET_SCALE_CONV, + MODEL_TENSOR.A_ENC_ATTN_SCALE, + MODEL_TENSOR.A_ENC_FFN_SCALE_LS, + MODEL_TENSOR.A_ENC_SPEAKER_PROJ, + MODEL_TENSOR.A_GEN_FLOW_INPUT_PROJ, + MODEL_TENSOR.A_GEN_FLOW_COND_EMBD, + MODEL_TENSOR.A_GEN_FLOW_TIME_FREQS, + MODEL_TENSOR.A_GEN_FLOW_TIME_UP, + MODEL_TENSOR.A_GEN_FLOW_TIME_DOWN, + MODEL_TENSOR.A_GEN_FLOW_TIME_NORM, + MODEL_TENSOR.A_GEN_FLOW_BLK_NORM, + MODEL_TENSOR.A_GEN_FLOW_BLK_UP, + MODEL_TENSOR.A_GEN_FLOW_BLK_DOWN, + MODEL_TENSOR.A_GEN_FLOW_BLK_ADA, + MODEL_TENSOR.A_GEN_FLOW_FINAL_ADA, + MODEL_TENSOR.A_GEN_FLOW_FINAL_PROJ, + MODEL_TENSOR.A_GEN_OUT_EOS, + MODEL_TENSOR.A_GEN_INPUT_LINEAR, + MODEL_TENSOR.A_GEN_EMB_MEAN, + MODEL_TENSOR.A_GEN_EMB_STD, + MODEL_TENSOR.A_GEN_WAV_QUANT_OUT, + MODEL_TENSOR.A_GEN_WAV_UPSAMPLE, + MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_IN, + MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_OUT, + MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV1, + MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV2, + MODEL_TENSOR.A_GEN_WAV_SEANET_SCALE_CONV, MODEL_TENSOR.A_ENC_CONV_NORM_MEAN, MODEL_TENSOR.A_ENC_CONV_NORM_VAR, MODEL_TENSOR.A_ENC_MEL_FILTERS, @@ -4628,6 +4726,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.D2T, ], MODEL_ARCH.DFLASH: [ + MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.ATTN_NORM, MODEL_TENSOR.ATTN_Q, @@ -4903,6 +5002,18 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_DOWN, MODEL_TENSOR.FFN_UP, ], + MODEL_ARCH.POCKETTTS: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], } # tensors that will not be serialized @@ -5179,6 +5290,8 @@ class VisionProjectorType: NEMOTRON_V2_VL = "nemotron_v2_vl" QWEN3TTS_SPKENC = "qwen3tts_spkenc" # audio: ECAPA-TDNN speaker encoder QWEN3TTS_GEN = "qwen3tts_gen" # audio generation: code_predictor + POCKETTTS_SPKENC = "pockettts_spkenc" # audio: mimi encoder as voice-prompt encoder + POCKETTTS_GEN = "pockettts_gen" # audio generation: flow-matching decoder + mimi decoder HUNYUANVL = "hunyuanvl" PARAKEET = "parakeet" # audio MINIMAXM3 = "minimax_m3" diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 81ae07c11a60..05f86396dc0a 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1453,6 +1453,9 @@ def add_gen_audio_head_count_kv(self, value: int) -> None: def add_gen_audio_attention_layernorm_eps(self, value: float) -> None: self.add_float32(Keys.ClipGenAudio.Attention.LAYERNORM_EPS, value) + def add_gen_audio_model_variant(self, value: str) -> None: + self.add_string(Keys.ClipGenAudio.MODEL_VARIANT, value) + def add_xielu_alpha_p(self, values: Sequence[float]): self.add_array(Keys.xIELU.ALPHA_P, values) diff --git a/include/llama.h b/include/llama.h index c14eb6f50d9a..bfef0e1d1129 100644 --- a/include/llama.h +++ b/include/llama.h @@ -203,11 +203,12 @@ extern "C" { }; enum llama_load_mode { - LLAMA_LOAD_MODE_NONE = 0, // no special loading mode - LLAMA_LOAD_MODE_MMAP = 1, // memory map the model - LLAMA_LOAD_MODE_MLOCK = 2, // force system to keep model in RAM rather than swapping or compressing - LLAMA_LOAD_MODE_MMAP_MLOCK = 3, // mmap + force system to keep model in RAM rather than swapping or compressing - LLAMA_LOAD_MODE_DIRECT_IO = 4, // use direct I/O if available + LLAMA_LOAD_MODE_AUTO = -1, // auto-detect based on device capabilities + LLAMA_LOAD_MODE_NONE = 0, // no special loading mode + LLAMA_LOAD_MODE_MMAP = 1, // memory map the model + LLAMA_LOAD_MODE_MLOCK = 2, // force system to keep model in RAM rather than swapping or compressing + LLAMA_LOAD_MODE_MMAP_MLOCK = 3, // mmap + force system to keep model in RAM rather than swapping or compressing + LLAMA_LOAD_MODE_DIRECT_IO = 4, // use direct I/O if available }; LLAMA_API const char * llama_load_mode_name(enum llama_load_mode load_mode); @@ -1324,8 +1325,6 @@ extern "C" { LLAMA_API void llama_sampler_apply ( struct llama_sampler * smpl, llama_token_data_array * cur_p); LLAMA_API void llama_sampler_reset ( struct llama_sampler * smpl); LLAMA_API struct llama_sampler * llama_sampler_clone (const struct llama_sampler * smpl); - // copy mutable sampler state without changing dst or its sampling graph bindings - // src and dst must have the same type and configuration LLAMA_API void llama_sampler_copy (const struct llama_sampler * src, struct llama_sampler * dst); // important: do not free if the sampler has been added to a llama_sampler_chain (via llama_sampler_chain_add) LLAMA_API void llama_sampler_free ( struct llama_sampler * smpl); diff --git a/requirements/requirements-convert_hf_to_gguf.txt b/requirements/requirements-convert_hf_to_gguf.txt index f80fdc1f6402..b1f7c863e27e 100644 --- a/requirements/requirements-convert_hf_to_gguf.txt +++ b/requirements/requirements-convert_hf_to_gguf.txt @@ -2,8 +2,4 @@ --extra-index-url https://download.pytorch.org/whl/cpu ## Embedding Gemma requires PyTorch 2.6.0 or later, bumped to 2.11.0 for compatibility -torch==2.11.0; platform_machine != "s390x" - -# torch s390x packages can only be found from nightly builds ---extra-index-url https://download.pytorch.org/whl/nightly -torch>=0.0.0.dev0; platform_machine == "s390x" +torch==2.11.0 diff --git a/requirements/requirements-convert_lora_to_gguf.txt b/requirements/requirements-convert_lora_to_gguf.txt index d091d564846b..5758076c41dc 100644 --- a/requirements/requirements-convert_lora_to_gguf.txt +++ b/requirements/requirements-convert_lora_to_gguf.txt @@ -1,4 +1,2 @@ -r ./requirements-convert_hf_to_gguf.txt --extra-index-url https://download.pytorch.org/whl/cpu -# torch s390x packages can only be found from nightly builds ---extra-index-url https://download.pytorch.org/whl/nightly diff --git a/scripts/fetch_server_test_models.py b/scripts/fetch_server_test_models.py deleted file mode 100755 index f43d1f63cdc7..000000000000 --- a/scripts/fetch_server_test_models.py +++ /dev/null @@ -1,105 +0,0 @@ -#!/usr/bin/env python -''' - This script fetches all the models used in the server tests. - - This is useful for slow tests that use larger models, to avoid them timing out on the model downloads. - - It is meant to be run from the root of the repository. - - Example: - python scripts/fetch_server_test_models.py - ( cd tools/server/tests && ./tests.sh -v -x -m slow ) -''' -import ast -import glob -import logging -import os -from typing import Generator -from pydantic import BaseModel -from typing import Optional -import subprocess - - -class HuggingFaceModel(BaseModel): - hf_repo: str - hf_file: Optional[str] = None - - class Config: - frozen = True - - -def collect_hf_model_test_parameters(test_file) -> Generator[HuggingFaceModel, None, None]: - try: - with open(test_file) as f: - tree = ast.parse(f.read()) - except Exception as e: - logging.error(f'collect_hf_model_test_parameters failed on {test_file}: {e}') - return - - for node in ast.walk(tree): - if isinstance(node, ast.FunctionDef): - for dec in node.decorator_list: - if isinstance(dec, ast.Call) and isinstance(dec.func, ast.Attribute) and dec.func.attr == 'parametrize': - param_names = ast.literal_eval(dec.args[0]).split(",") - if "hf_repo" not in param_names: - continue - - raw_param_values = dec.args[1] - if not isinstance(raw_param_values, ast.List): - logging.warning(f'Skipping non-list parametrize entry at {test_file}:{node.lineno}') - continue - - hf_repo_idx = param_names.index("hf_repo") - hf_file_idx = param_names.index("hf_file") if "hf_file" in param_names else None - - for t in raw_param_values.elts: - if not isinstance(t, ast.Tuple): - logging.warning(f'Skipping non-tuple parametrize entry at {test_file}:{node.lineno}') - continue - yield HuggingFaceModel( - hf_repo=ast.literal_eval(t.elts[hf_repo_idx]), - hf_file=ast.literal_eval(t.elts[hf_file_idx]) if hf_file_idx is not None else None) - - -if __name__ == '__main__': - logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s') - - models = sorted(list(set([ - model - for test_file in glob.glob('tools/server/tests/unit/test_*.py') - for model in collect_hf_model_test_parameters(test_file) - ])), key=lambda m: (m.hf_repo, m.hf_file)) - - logging.info(f'Found {len(models)} models in parameterized tests:') - for m in models: - logging.info(f' - {m.hf_repo} / {m.hf_file}') - - cli_path = os.environ.get( - 'LLAMA_CLI_BIN_PATH', - os.path.join( - os.path.dirname(__file__), - '../build/bin/Release/llama-cli.exe' if os.name == 'nt' else '../build/bin/llama-cli')) - - for m in models: - if '<' in m.hf_repo or (m.hf_file is not None and '<' in m.hf_file): - continue - if m.hf_file is not None and '-of-' in m.hf_file: - logging.warning(f'Skipping model at {m.hf_repo} / {m.hf_file} because it is a split file') - continue - logging.info(f'Using llama-cli to ensure model {m.hf_repo}/{m.hf_file} was fetched') - cmd = [ - cli_path, - '-hfr', m.hf_repo, - *([] if m.hf_file is None else ['-hff', m.hf_file]), - '-n', '1', - '-p', 'Hey', - '--no-warmup', - '--log-disable', - '-st'] - if m.hf_file != 'tinyllamas/stories260K.gguf' and 'Mistral-Nemo' not in m.hf_repo: - cmd += ('-fa', 'on') - try: - subprocess.check_call(cmd) - except subprocess.CalledProcessError: - logging.error(f'Failed to fetch model at {m.hf_repo} / {m.hf_file} with command:\n {" ".join(cmd)}') - exit(1) diff --git a/scripts/hip/gcn-cdna-vgpr-check.py b/scripts/hip/gcn-cdna-vgpr-check.py index bbbce52ef39f..40fb789417c3 100644 --- a/scripts/hip/gcn-cdna-vgpr-check.py +++ b/scripts/hip/gcn-cdna-vgpr-check.py @@ -60,90 +60,10 @@ def main(): log_file = sys.argv[1] ignored = { '_ZL21gated_linear_attn_f32ILi128EEviiiifPKfS1_S1_S1_S1_Pf', - '_ZL18flash_attn_ext_f16ILi64ELi64ELi16ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi80ELi80ELi16ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi96ELi96ELi16ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi64ELi64ELi32ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', '_ZL13rwkv_wkv7_f32ILi128EEviiiiPKfS1_S1_S1_S1_S1_S1_Pf', - '_ZL18flash_attn_ext_f16ILi80ELi80ELi16ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi16ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi80ELi80ELi32ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi96ELi96ELi16ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi16ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi16ELi2ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi96ELi96ELi32ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi16ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi32ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi16ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi16ELi1ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi80ELi80ELi2ELi8ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi96ELi96ELi2ELi8ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi2ELi8ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi2ELi8ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi2ELi8ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi16ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi16ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi16ELi4ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi32ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi96ELi96ELi4ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi4ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi4ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi4ELi4ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi80ELi80ELi4ELi8ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi4ELi8ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi96ELi96ELi64ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi64ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi64ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi64ELi1ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi64ELi64ELi8ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi80ELi80ELi8ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi96ELi96ELi8ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi8ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi80ELi80ELi8ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi8ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi8ELi4ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi96ELi96ELi8ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi8ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi8ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi8ELi2ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi8ELi8ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi8ELi8ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi8ELi8ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL24mul_mat_q_stream_k_fixupIL9ggml_type22ELi8ELb1EEvPKiS2_PfPKfiiimimimi', - '_ZL9mul_mat_qIL9ggml_type3ELi32ELb0EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type3ELi48ELb0EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type20ELi32ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type17ELi64ELb0EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL18flash_attn_ext_f16ILi80ELi80ELi4ELi4ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL15flash_attn_tileILi256ELi256ELi32ELi1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL9mul_mat_qIL9ggml_type19ELi112ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type17ELi112ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type22ELi112ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type19ELi128ELb0EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type19ELi128ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type7ELi112ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type3ELi128ELb0EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type3ELi128ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type7ELi128ELb0EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type7ELi128ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type11ELi112ELb0EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type11ELi112ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL24mul_mat_q_stream_k_fixupIL9ggml_type11ELi128ELb0EEvPKiS2_PfPKfiiimimimi', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi32ELi1ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL9mul_mat_qIL9ggml_type2ELi112ELb0EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi32ELi2ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi112ELi112ELi4ELi8ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi32ELi1ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi32ELi2ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi128ELi128ELi4ELi8ELb1ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_f16ILi96ELi96ELi4ELi8ELb0ELb0EEvPKcS1_S1_S1_S1_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS5_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL18flash_attn_ext_vecILi128ELi2EL9ggml_type2ELS0_2ELb0EEvPKcS2_S2_S2_S2_PKiPfP15HIP_vector_typeIfLj2EEffffjfiS6_IjLj3EEiiiiiiiiiiiliiliiiiil', - '_ZL9mul_mat_qIL9ggml_type10ELi16ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type12ELi128ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type40ELi112ELb0EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type40ELi112ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type40ELi128ELb0EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii', - '_ZL9mul_mat_qIL9ggml_type40ELi128ELb1EEvPKcPKiS4_S4_PfS5_iiiiiiiiiiiiiiiii' + '_ZL12rwkv_wkv_f32ILi128EEviiiiPKfS1_S1_S1_S1_S1_Pf', + '_ZL9mul_mat_qIL9ggml_type10ELi64ELb1EEvPKcPKiS4_S4_PfS5_PKf15HIP_vector_typeIjLj3EEiiiiiS9_S9_iiiS9_S9_iiiS9_', + '_ZL9mul_mat_qIL9ggml_type42ELi128ELb1EEvPKcPKiS4_S4_PfS5_PKf15HIP_vector_typeIjLj3EEiiiiiS9_S9_iiiS9_S9_iiiS9_', } functions = parse_log_file(log_file) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 73fb8b981382..8ed9391d7c74 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -147,6 +147,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_MELLUM, "mellum" }, { LLM_ARCH_NANBEIGE, "nanbeige" }, { LLM_ARCH_QWEN3TTS, "qwen3tts" }, + { LLM_ARCH_POCKETTTS, "pockettts" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; diff --git a/src/llama-arch.h b/src/llama-arch.h index 51dfd288ced2..18d9de186f75 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -152,6 +152,7 @@ enum llm_arch { LLM_ARCH_DFLASH, LLM_ARCH_NANBEIGE, LLM_ARCH_QWEN3TTS, + LLM_ARCH_POCKETTTS, LLM_ARCH_UNKNOWN, }; diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 3d50f8a1cb1c..51ba05439682 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -543,7 +543,7 @@ llama_model_loader::llama_model_loader( tensor_buft_overrides = param_tensor_buft_overrides_p; - this->use_mmap = load_mode == LLAMA_LOAD_MODE_MMAP || load_mode == LLAMA_LOAD_MODE_MMAP_MLOCK; + this->use_mmap = load_mode == LLAMA_LOAD_MODE_MMAP || load_mode == LLAMA_LOAD_MODE_MMAP_MLOCK || load_mode == LLAMA_LOAD_MODE_AUTO; this->use_direct_io = load_mode == LLAMA_LOAD_MODE_DIRECT_IO; if (!fname.empty()) { diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 3bf3a22f26af..0e27cb41713b 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -116,6 +116,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_qwen3vlmoe(params); case LLM_ARCH_QWEN3TTS: return new llama_model_qwen3tts(params); + case LLM_ARCH_POCKETTTS: + return new llama_model_pockettts(params); case LLM_ARCH_PHI2: return new llama_model_phi2(params); case LLM_ARCH_PHI3: @@ -1271,8 +1273,23 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { this->ml = &ml; // to be used by create_tensor() and load_arch_tensors() + if (ml.use_mmap && params.load_mode == LLAMA_LOAD_MODE_AUTO) { + for (const auto & dev : devices) { + ggml_backend_dev_props props; + ggml_backend_dev_get_props(dev.dev, &props); + if (!props.caps.mmap_support) { + ml.use_mmap = false; + break; + } + } + } + + const char * load_mode_name = params.load_mode == LLAMA_LOAD_MODE_AUTO + ? llama_load_mode_name(ml.use_mmap ? LLAMA_LOAD_MODE_MMAP : LLAMA_LOAD_MODE_NONE) + : llama_load_mode_name(params.load_mode); + LLAMA_LOG_INFO("%s: loading model tensors, this can take a while... (load_mode = %s)\n", - __func__, llama_load_mode_name(params.load_mode)); + __func__, load_mode_name); // build a list of buffer types for the CPU and GPU devices pimpl->cpu_buft_list = make_cpu_buft_list(devices, params.use_extra_bufts, params.no_host); @@ -2452,7 +2469,7 @@ llama_model_params llama_model_default_params() { /*.tensor_buft_overrides =*/ nullptr, /*.n_gpu_layers =*/ -1, /*.split_mode =*/ LLAMA_SPLIT_MODE_LAYER, - /*.load_mode =*/ LLAMA_LOAD_MODE_MMAP, + /*.load_mode =*/ LLAMA_LOAD_MODE_AUTO, /*.main_gpu =*/ 0, /*.tensor_split =*/ nullptr, /*.progress_callback =*/ nullptr, @@ -2622,6 +2639,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_MAINCODER: case LLM_ARCH_GLM_DSA: case LLM_ARCH_NANBEIGE: + case LLM_ARCH_POCKETTTS: return LLAMA_ROPE_TYPE_NORM; // the pairs of head values are offset by n_rot/2 diff --git a/src/llama-model.h b/src/llama-model.h index 1dd0904387af..341cb66fbafe 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -623,8 +623,9 @@ struct llama_model { struct ggml_tensor * per_layer_model_proj = nullptr; struct ggml_tensor * per_layer_proj_norm = nullptr; - // eagle3 - struct ggml_tensor * fc = nullptr; // feature fusion layer + // eagle3 / dflash feature fusion layer + struct ggml_tensor * fc = nullptr; + struct ggml_tensor * fc_s = nullptr; struct ggml_tensor * d2t = nullptr; // draft to target vocabulary mapping // dspark diff --git a/src/llama.cpp b/src/llama.cpp index d6e0bbfefa72..94c8f60e0c43 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -48,6 +48,8 @@ const char * llama_flash_attn_type_name(enum llama_flash_attn_type flash_attn_ty const char * llama_load_mode_name(enum llama_load_mode load_mode) { switch (load_mode) { + case LLAMA_LOAD_MODE_AUTO: + return "auto"; case LLAMA_LOAD_MODE_NONE: return "none"; case LLAMA_LOAD_MODE_MMAP: @@ -63,11 +65,12 @@ const char * llama_load_mode_name(enum llama_load_mode load_mode) { } enum llama_load_mode llama_load_mode_from_str(const char * str) { - if (std::strcmp(str, "none") == 0) { return LLAMA_LOAD_MODE_NONE; } - if (std::strcmp(str, "mmap") == 0) { return LLAMA_LOAD_MODE_MMAP; } - if (std::strcmp(str, "mlock") == 0) { return LLAMA_LOAD_MODE_MLOCK; } + if (std::strcmp(str, "auto") == 0) { return LLAMA_LOAD_MODE_AUTO; } + if (std::strcmp(str, "none") == 0) { return LLAMA_LOAD_MODE_NONE; } + if (std::strcmp(str, "mmap") == 0) { return LLAMA_LOAD_MODE_MMAP; } + if (std::strcmp(str, "mlock") == 0) { return LLAMA_LOAD_MODE_MLOCK; } if (std::strcmp(str, "mmap+mlock") == 0) { return LLAMA_LOAD_MODE_MMAP_MLOCK; } - if (std::strcmp(str, "dio") == 0) { return LLAMA_LOAD_MODE_DIRECT_IO; } + if (std::strcmp(str, "dio") == 0) { return LLAMA_LOAD_MODE_DIRECT_IO; } throw std::invalid_argument(std::string("unknown load mode: ") + str); } diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index daff6e78f1cf..bfbdb28ee5e9 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -79,6 +79,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { const int64_t n_embd_inp = hparams.n_embd_inp_enc(); + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED); // DSpark = DFlash + a semi-autoregressive Markov head and Confidence head // // TODO: only Qwen3-style backbones are supported for now; other backbones (e.g. Gemma4) @@ -97,6 +98,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { } 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) output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); // decoder final norm @@ -205,7 +207,7 @@ template <> llama_model_dflash::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { ggml_tensor * cur = build_inp_embd_enc(); - cur = build_lora_mm(model.fc, cur); + cur = build_lora_mm(model.fc, cur, model.fc_s); cb(cur, "fc_out", -1); cur = build_norm(cur, model.output_norm_enc, NULL, LLM_NORM_RMS, -1); @@ -460,9 +462,9 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra cb(cur, "ffn_norm", il); cur = build_ffn(cur, - layer.ffn_up, NULL, NULL, - layer.ffn_gate, NULL, NULL, - layer.ffn_down, NULL, NULL, + layer.ffn_up, NULL, layer.ffn_up_s, + layer.ffn_gate, NULL, layer.ffn_gate_s, + layer.ffn_down, NULL, layer.ffn_down_s, NULL, LLM_FFN_SILU, LLM_FFN_PAR, il); cb(cur, "ffn_out", il); @@ -479,15 +481,17 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra res->t_embd = cur; // lm_head from the target model (shared via ctx_other) - auto * output = model.output; + auto * output = model.output; + auto * output_s = model.output_s; if (output == nullptr) { GGML_ASSERT(cparams.ctx_other != nullptr); const auto * model_other = llama_get_model(cparams.ctx_other); GGML_ASSERT(model_other->output != nullptr && "DFlash decoder requires the target model's output projection"); - output = model_other->output; + output = model_other->output; + output_s = model_other->output_s; } - cur = build_lora_mm(output, cur); + cur = build_lora_mm(output, cur, output_s); cb(cur, "result_output", -1); res->t_logits = cur; @@ -655,15 +659,17 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_ cb(cur, "result_norm", -1); // lm_head from the target model (shared via ctx_other) - auto * output = model.output; + auto * output = model.output; + auto * output_s = model.output_s; if (output == nullptr) { GGML_ASSERT(cparams.ctx_other != nullptr); const auto * model_other = llama_get_model(cparams.ctx_other); GGML_ASSERT(model_other->output != nullptr && "DSpark decoder requires the target model's output projection"); - output = model_other->output; + output = model_other->output; + output_s = model_other->output_s; } - cur = build_lora_mm(output, cur); + cur = build_lora_mm(output, cur, output_s); cb(cur, "result_output", -1); res->t_logits = cur; diff --git a/src/models/exaone4.cpp b/src/models/exaone4.cpp index 863268abcef2..a06819a67caa 100644 --- a/src/models/exaone4.cpp +++ b/src/models/exaone4.cpp @@ -1,6 +1,9 @@ #include "models.h" void llama_model_exaone4::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); + GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer"); + if (hparams.n_layer() == 64) { // 32B hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; hparams.n_swa = 4096; @@ -15,9 +18,6 @@ void llama_model_exaone4::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false); ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); - ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); - - GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer"); switch (hparams.n_layer()) { case 30: type = LLM_TYPE_1_2B; break; diff --git a/src/models/models.h b/src/models/models.h index 9230345208a9..ddb9ae2f1210 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -713,6 +713,19 @@ struct llama_model_gpt2 : public llama_model_base { }; +struct llama_model_pockettts : public llama_model_base { + llama_model_pockettts(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + struct graph : public llm_graph_context { + graph(const llama_model & model, const llm_graph_params & params); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_codeshell : public llama_model_base { llama_model_codeshell(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/src/models/nemotron-h.cpp b/src/models/nemotron-h.cpp index cd2af3179c3a..f02674c64610 100644 --- a/src/models/nemotron-h.cpp +++ b/src/models/nemotron-h.cpp @@ -177,8 +177,11 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_ auto * inp = build_inp_mem_hybrid(); ggml_tensor * inp_out_ids = build_inp_out_ids(); + const bool extract_final_inp = (size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer]; for (int il = 0; il < n_layer; ++il) { + res->t_layer_inp[il] = inpL; + struct ggml_tensor * inpSA = inpL; // norm @@ -195,7 +198,7 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_ cur = build_ffn_layer(cur, model, il); } - if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) { + if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked && !extract_final_inp) { cur = ggml_get_rows(ctx0, cur, inp_out_ids); inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); } @@ -209,6 +212,13 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_ } cur = inpL; + if (extract_final_inp) { + res->t_layer_inp[n_layer] = cur; + + if (inp_out_ids && cparams.embeddings_nextn_masked) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + } cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); diff --git a/src/models/pockettts.cpp b/src/models/pockettts.cpp new file mode 100644 index 000000000000..1b3bb6c648af --- /dev/null +++ b/src/models/pockettts.cpp @@ -0,0 +1,146 @@ +#include "models.h" + +// backbone of the pocket-tts CALM pipeline: the "text" side of a flow language model. +// it has no lm_head, the audio latents are produced by the flow net inside the mmproj + +void llama_model_pockettts::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps); + + switch (hparams.n_layer()) { + case 6: type = LLM_TYPE_109M; break; + case 24: type = LLM_TYPE_335M; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_pockettts::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output_norm_b = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "bias"), {n_embd}, 0); + // no output head, the logits are unused; reuse the embedding table so a sampler can still run + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + layer.attn_norm_b = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "bias", i), {n_embd}, 0); + + create_tensor_qkv(layer, i, n_embd, n_embd, n_embd_gqa, n_embd_gqa, TENSOR_NOT_REQUIRED); + + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd, n_embd}, 0); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_norm_b = create_tensor(tn(LLM_TENSOR_FFN_NORM, "bias", i), {n_embd}, 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); + } +} + +std::unique_ptr llama_model_pockettts::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +llama_model_pockettts::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { + const int64_t n_embd_head = hparams.n_embd_head_v(); + + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + GGML_ASSERT(n_embd_head == n_rot); + + ggml_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + + ggml_tensor * inp_pos = build_inp_pos(); + + auto * inp_attn = build_attn_inp_kv(); + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + for (int il = 0; il < n_layer; ++il) { + cur = build_norm(inpL, + model.layers[il].attn_norm, + model.layers[il].attn_norm_b, + LLM_NORM, il); + cb(cur, "attn_norm", il); + + // self-attention + { + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head, n_head, n_head_kv, 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 + ); + + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + + cur = build_attn(inp_attn, + model.layers[il].wo, NULL, model.layers[il].wo_s, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il); + } + + if (il == n_layer - 1 && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inpL = ggml_get_rows(ctx0, inpL, inp_out_ids); + } + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL); + cb(ffn_inp, "ffn_inp", il); + + // FF + { + cur = build_norm(ffn_inp, + model.layers[il].ffn_norm, + model.layers[il].ffn_norm_b, + LLM_NORM, il); + cb(cur, "ffn_norm", il); + + cur = build_ffn(cur, + model.layers[il].ffn_up, NULL, NULL, + NULL, NULL, NULL, + model.layers[il].ffn_down, NULL, NULL, + NULL, + LLM_FFN_GELU, LLM_FFN_SEQ, il); + cb(cur, "ffn_out", il); + } + + cur = ggml_add(ctx0, cur, ffn_inp); + + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + + // input for next layer + inpL = cur; + } + + cur = build_norm(inpL, + model.output_norm, + model.output_norm_b, + LLM_NORM, -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); +} diff --git a/tests/peg-parser/test-json-parser.cpp b/tests/peg-parser/test-json-parser.cpp index 5dd00115cead..ec7c2e668ff1 100644 --- a/tests/peg-parser/test-json-parser.cpp +++ b/tests/peg-parser/test-json-parser.cpp @@ -77,6 +77,30 @@ void test_json_parser(testing &t) { t.assert_equal("result_is_need_more_input", true, result.need_more_input()); }); + // Test need_more_input() parsing - incomplete escape sequence in a string value + t.test("need_more_input() parsing - incomplete escape sequence", [](testing &t) { + auto json = build_peg_parser([](common_peg_parser_builder & p) { return p.json(); }); + + std::vector inputs { + R"({"text": "hello\)", // dangling backslash + R"({"text": "hello\u)", // incomplete unicode escape sequence + R"({"text": "hello\u00)", + }; + + for (const auto & input : inputs) { + t.test(input, [&](testing &t) { + common_peg_parse_context ctx(input, COMMON_PEG_PARSE_FLAG_LENIENT); + + auto result = json.parse(ctx); + + t.assert_equal("result_is_need_more_input", true, result.need_more_input()); + + // the incomplete escape sequence is not part of the partial value + t.assert_equal("result_end", input.find('\\'), result.end); + }); + } + }); + t.test("object member", [](testing &t) { auto parser = build_peg_parser([](common_peg_parser_builder & p) { return p.json_member("name", "\"" + p.chars("[a-z]") + "\""); diff --git a/tests/test-backend-sampler.cpp b/tests/test-backend-sampler.cpp index 4b3d12635a36..c23e7248d5e5 100644 --- a/tests/test-backend-sampler.cpp +++ b/tests/test-backend-sampler.cpp @@ -2105,7 +2105,9 @@ static std::vector collect_tests_to_run(const std::st #ifdef GGML_USE_HIP // TODO: remove this when https://github.com/ggml-org/llama.cpp/pull/26592 is merged if (test.name == "penalties" || test.name == "set_sampler" || - test.name == "mixed" || test.name == "top_p") { + test.name == "mixed" || test.name == "top_p" || + test.name == "multi_output_sampling_chain" || + test.name == "multi_output_cpu") { fprintf(stderr, "Skipping test '%s' on HIP backend (no backend TOP_K support)\n", test.name.c_str()); continue; } diff --git a/tools/cli/README.md b/tools/cli/README.md index b42b2e5343ff..880c4a540838 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -58,7 +58,7 @@ | `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing
(env: LLAMA_ARG_MLOCK) | | `--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: mmap)
- 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) | +| `-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) | | `--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 | diff --git a/tools/completion/README.md b/tools/completion/README.md index 552a0c6abf4e..c2e52ac066e7 100644 --- a/tools/completion/README.md +++ b/tools/completion/README.md @@ -141,7 +141,7 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1 | `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing
(env: LLAMA_ARG_MLOCK) | | `--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: mmap)
- 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) | +| `-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) | | `--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 | diff --git a/tools/imatrix/imatrix.cpp b/tools/imatrix/imatrix.cpp index 3431a4eca84b..f5fee621840a 100644 --- a/tools/imatrix/imatrix.cpp +++ b/tools/imatrix/imatrix.cpp @@ -222,6 +222,15 @@ static void compute_cossim(std::vector & tstats) { } } +static bool all_finite(const float * v, size_t n) { + for (size_t i = 0; i < n; ++i) { + if (!std::isfinite(v[i])) { + return false; + } + } + return true; +} + bool IMatrixCollector::collect_imatrix(struct ggml_tensor * t, bool ask, void * user_data) { GGML_UNUSED(user_data); @@ -299,33 +308,39 @@ bool IMatrixCollector::collect_imatrix(struct ggml_tensor * t, bool ask, void * exit(1); //GGML_ABORT("fatal error"); } LOG_DBGV(2, "%s[%d]: %32s, %s, %5d x %5d, %d\n", __func__, m_last_chunk, wname.c_str(), ggml_op_name(t->op), (int)src1->ne[0], (int)src1->ne[2], (int)src1->type); - // loop over all possible experts, regardless if they are used or not in the batch - for (int64_t ex = 0; ex < n_as; ++ex) { - size_t e_start = ex*src1->ne[0]; - - for (int64_t idx = 0; idx < n_ids; ++idx) { - for (int64_t row = 0; row < src1->ne[2]; ++row) { - const int excur = *(const int32_t *) (m_ids.data() + row*ids->nb[1] + idx*ids->nb[0]); - GGML_ASSERT(excur >= 0 && excur < n_as); // sanity check + const int64_t ne0 = src1->ne[0]; + const int64_t n_tokens = src1->ne[2]; - if (excur != ex) continue; + // single pass over the routing ids + std::vector touched(n_as, 0); + for (int64_t idx = 0; idx < n_ids; ++idx) { + for (int64_t row = 0; row < n_tokens; ++row) { + const int32_t ex = *(const int32_t *) (m_ids.data() + row * ids->nb[1] + idx * ids->nb[0]); - const int64_t i11 = idx % src1->ne[1]; - const int64_t i12 = row; - const float * x = (const float *)(data + i11*src1->nb[1] + i12*src1->nb[2]); + GGML_ASSERT(ex >= 0 && ex < n_as); // sanity check - e.counts[ex]++; + const int64_t i11 = idx % src1->ne[1]; + const float * x = (const float *) (data + i11 * src1->nb[1] + row * src1->nb[2]); + float * acc = e.values.data() + ex * ne0; - for (int64_t j = 0; j < src1->ne[0]; ++j) { - e.values[e_start + j] += x[j] * x[j]; - if (!std::isfinite((float)e.values[e_start + j])) { - LOG_ERR("%f detected in %s\n", (float)e.values[e_start + j], wname.c_str()); - exit(1); - } - } + e.counts[ex]++; + touched[ex] = 1; + for (int64_t j = 0; j < ne0; ++j) { + acc[j] += x[j] * x[j]; } } + } + + // check for non-finite values, only checking experts that were routed to and touched + for (int64_t ex = 0; ex < n_as; ++ex) { + if (touched[ex] && !all_finite(e.values.data() + ex * ne0, ne0)) { + LOG_ERR("%s: non-finite values detected in %s\n", __func__, wname.c_str()); + exit(1); + } + } + + for (int64_t ex = 0; ex < n_as; ++ex) { const int32_t n_chunk = e.counts[ex] / chunk_size; if (n_chunk > m_last_chunk) { const int32_t chunk_step = n_chunk - m_last_chunk; @@ -366,24 +381,28 @@ bool IMatrixCollector::collect_imatrix(struct ggml_tensor * t, bool ask, void * } LOG_DBGV(2, "%s[%d]: %32s, %s, %5d x %5d x %5d, %d\n", __func__, m_last_chunk, wname.c_str(), ggml_op_name(t->op), (int)src1->ne[0], (int)src1->ne[1], (int)src1->ne[2], (int)src1->type); + const int64_t ne0 = src1->ne[0]; + for (int64_t i3 = 0; i3 < src1->ne[3]; ++i3) { for (int64_t i2 = 0; i2 < src1->ne[2]; ++i2) { // handle 3D+ tensors, but flatten 3D+ activations when model tensor is 2D const int64_t mat_id = (i3 % src0->ne[3]) * src0->ne[2] + (i2 % src0->ne[2]); - const int64_t mat_start = mat_id * src1->ne[0]; + float * acc = e.values.data() + mat_id * ne0; for (int64_t row = 0; row < src1->ne[1]; ++row) { const float * x = (const float *) (data + row * src1->nb[1] + i2 * src1->nb[2] + i3 * src1->nb[3]); - for (int64_t j = 0; j < src1->ne[0]; ++j) { - e.values[mat_start + j] += x[j] * x[j]; - if (!std::isfinite((float)e.values[j])) { - LOG_ERR("%f detected in %s\n", (float)e.values[j], wname.c_str()); - exit(1); - } + for (int64_t j = 0; j < ne0; ++j) { + acc[j] += x[j] * x[j]; } } } } + + // check for non-finite values + if (!all_finite(e.values.data(), e.values.size())) { + LOG_ERR("%s: non-finite values detected in %s\n", __func__, wname.c_str()); + exit(1); + } // only 1 count in practice, except when a tensor is used for both MUL_MAT_ID and MUL_MAT for (size_t i = 0; i < e.counts.size(); ++i) { e.counts[i] += ggml_nrows(src1) / n_mat; diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index c17a27b54019..7c495afe20ef 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -384,7 +384,7 @@ static const cmd_params cmd_params_defaults = { /* n_gpu_layers */ { -1 }, /* n_cpu_moe */ { 0 }, /* split_mode */ { LLAMA_SPLIT_MODE_LAYER }, - /* load_mode */ { LLAMA_LOAD_MODE_MMAP }, + /* load_mode */ { LLAMA_LOAD_MODE_AUTO }, /* main_gpu */ { 0 }, /* no_kv_offload */ { false }, /* flash_attn */ { LLAMA_FLASH_ATTN_TYPE_AUTO }, @@ -459,7 +459,7 @@ static void print_usage(int /* argc */, char ** argv) { printf(" -nkvo, --no-kv-offload <0|1> (default: %s)\n", join(cmd_params_defaults.no_kv_offload, ",").c_str()); 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(" -lm, --load-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").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()); @@ -764,7 +764,9 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { std::vector modes; for (const auto & m : p) { llama_load_mode mode; - if (m == "none") { + if (m == "auto") { + mode = LLAMA_LOAD_MODE_AUTO; + } else if (m == "none") { mode = LLAMA_LOAD_MODE_NONE; } else if (m == "mmap") { mode = LLAMA_LOAD_MODE_MMAP; diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index fe22cb12543e..3f4a6c670dd7 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -57,6 +57,9 @@ add_library(mtmd models/mimo-audio.cpp models/qwen3tts-spkenc.cpp models/qwen3tts-gen.cpp + models/pockettts-seanet.cpp + models/pockettts-spkenc.cpp + models/pockettts-gen.cpp models/step3vl.cpp models/siglip.cpp models/whisper-enc.cpp diff --git a/tools/mtmd/README-dev.md b/tools/mtmd/README-dev.md index 3cddd085ec6c..ac43e1b81b1c 100644 --- a/tools/mtmd/README-dev.md +++ b/tools/mtmd/README-dev.md @@ -59,8 +59,10 @@ Due to wide variety of audio generation pipelines, the `mtmd_gen_audio` system i ### Checklist for porting new audio generation models to mtmd -1. Establish a list of reusable and missing components from the current mtmd implementation. -2. For GGUF conversion: +1. Make sure to consult merged PRs about adding new TTS models, especially reviewer comments + - Example: https://github.com/ggml-org/llama.cpp/pulls?q=is%3Apr+mtmd+tts+is%3Amerged +2. Establish a list of reusable and missing components from the current mtmd implementation. +3. For GGUF conversion: - Backbone model should be converted to a normal text model (loadable via `libllama`) - If model used hard-coded embedding row ID, append them to token embeddings and assign token name for them (see `qwen3tts.py`) - If model have a specific output logits head for audio codes (usually semantic code), keep the head as-is and pad the logits at inference time (see `src/models/qwen3vl.cpp`) @@ -70,12 +72,17 @@ Due to wide variety of audio generation pipelines, the `mtmd_gen_audio` system i - For tensor naming: - Prefixed with `a.*` for tensors used by speaker encoder pipeline - Prefixed with `a.gen.*` for generation stages (code / mel-spectrogram / PCM generation) -3. Make sure most of the changes happen inside `mtmd-helper-gen.cpp`. A good PR looks like this: + - For GGUF metadata: + - Reuse as many existing keys as possible + - In most cases, you can hard-code model configs in the model graph class, or in `clip_hparams` + - If some values need to be exposed to the `mtmd_helper` layer, hard-code them in `mtmd_helper` and distinguish by pipeline and `mtmd_gen_audio_info::model_variant` if necessary + - Do NOT add new GGUF metadata or new fields to `mtmd_gen_audio_info` unless you can prove that you absolutely need them +4. Make sure most of the changes happen inside `mtmd-helper-gen.cpp`. A good PR looks like this: - 10-20% changes is to add new backbone (text) model and conversion - 60% changes inside `mtmd-helper-gen.cpp` - 10% changes inside `libmtmd` and `clip.cpp` systems - The rest downstream code (CLI, server) should have no changes at all -4. Update usage documentation in `tools/tts/README.md` +5. Update usage documentation in `tools/tts/README.md` IMPORTANT: If your model needs changes that don't fit the existing infrastructure, **open an issue first for discussion**. diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index bf73222f9def..b2c8b4021321 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -92,7 +92,9 @@ #define KEY_A_LOCAL_GROUP_SIZE "clip.audio.local_group_size" // mimo-v2.5: input_local_transformer grouping size // audio generation (gen-audio)-specific #define KEY_GEN_AUDIO_PROJ_TYPE "clip.gen.audio.projector_type" // for models with mixed modalities -#define KEY_AUDIO_SUBSAMPLING_FACTOR "clip.audio.subsampling_factor" +// name of the weight variant, for settings that are not in the checkpoint +#define KEY_GEN_AUDIO_VARIANT "clip.gen.audio.model_variant" +#define KEY_AUDIO_SUBSMPL_FACTOR "clip.audio.subsampling_factor" // // tensor name constants @@ -246,6 +248,38 @@ #define TN_A_GEN_WAV_DAC_POST_SNAKE "a.gen.wav.dac.post_snake.%s" #define TN_A_GEN_WAV_DAC_POST_CONV "a.gen.wav.dac.post_conv.%s" +// pocket-tts +#define TN_A_SEANET_CONV_IN "a.seanet.conv_in.%s" +#define TN_A_SEANET_CONV_OUT "a.seanet.conv_out.%s" +#define TN_A_SEANET_RES_CONV1 "a.seanet.blk.%d.res_conv1.%s" +#define TN_A_SEANET_RES_CONV2 "a.seanet.blk.%d.res_conv2.%s" +#define TN_A_SEANET_SCALE_CONV "a.seanet.blk.%d.scale_conv.%s" +#define TN_A_SPEAKER_PROJ "a.speaker_proj.%s" +#define TN_A_DOWNSAMPLE_CONV "a.downsample.conv.%s" +#define TN_A_GEN_FLOW_INPUT_PROJ "a.gen.flow.input_proj.%s" +#define TN_A_GEN_FLOW_COND_EMBD "a.gen.flow.cond_embd.%s" +#define TN_A_GEN_FLOW_TIME_FREQS "a.gen.flow.time.%d.freqs" +#define TN_A_GEN_FLOW_TIME_UP "a.gen.flow.time.%d.up.%s" +#define TN_A_GEN_FLOW_TIME_DOWN "a.gen.flow.time.%d.down.%s" +#define TN_A_GEN_FLOW_TIME_NORM "a.gen.flow.time.%d.norm" +#define TN_A_GEN_FLOW_BLK_NORM "a.gen.flow.blk.%d.norm.%s" +#define TN_A_GEN_FLOW_BLK_UP "a.gen.flow.blk.%d.up.%s" +#define TN_A_GEN_FLOW_BLK_DOWN "a.gen.flow.blk.%d.down.%s" +#define TN_A_GEN_FLOW_BLK_ADA "a.gen.flow.blk.%d.ada.%s" +#define TN_A_GEN_FLOW_FINAL_ADA "a.gen.flow.final.ada.%s" +#define TN_A_GEN_FLOW_FINAL_PROJ "a.gen.flow.final.proj.%s" +#define TN_A_GEN_OUT_EOS "a.gen.out_eos.%s" +#define TN_A_GEN_INPUT_LINEAR "a.gen.input_linear.%s" +#define TN_A_GEN_EMB_MEAN "a.gen.emb_mean" +#define TN_A_GEN_EMB_STD "a.gen.emb_std" +#define TN_A_GEN_WAV_QUANT_OUT "a.gen.wav.quant_out.%s" +#define TN_A_GEN_WAV_UPSAMPLE "a.gen.wav.upsample.%s" +#define TN_A_GEN_WAV_SEANET_CONV_IN "a.gen.wav.seanet.conv_in.%s" +#define TN_A_GEN_WAV_SEANET_CONV_OUT "a.gen.wav.seanet.conv_out.%s" +#define TN_A_GEN_WAV_SEANET_RES_CONV1 "a.gen.wav.seanet.blk.%d.res_conv1.%s" +#define TN_A_GEN_WAV_SEANET_RES_CONV2 "a.gen.wav.seanet.blk.%d.res_conv2.%s" +#define TN_A_GEN_WAV_SEANET_SCALE_CONV "a.gen.wav.seanet.blk.%d.scale_conv.%s" + // cogvlm #define TN_MM_POST_FC_NORM "mm.post_fc_norm.%s" #define TN_MM_H_TO_4H "mm.up.%s" @@ -455,6 +489,8 @@ enum projector_type { PROJECTOR_TYPE_MIMO_AUDIO, PROJECTOR_TYPE_QWEN3TTS_SPKENC, PROJECTOR_TYPE_QWEN3TTS_GEN, + PROJECTOR_TYPE_POCKETTTS_SPKENC, + PROJECTOR_TYPE_POCKETTTS_GEN, PROJECTOR_TYPE_MUSE_GLIMMER, PROJECTOR_TYPE_UNKNOWN, }; @@ -515,6 +551,8 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_PARAKEET, "parakeet"}, { PROJECTOR_TYPE_QWEN3TTS_SPKENC, "qwen3tts_spkenc"}, { PROJECTOR_TYPE_QWEN3TTS_GEN, "qwen3tts_gen"}, + { PROJECTOR_TYPE_POCKETTTS_SPKENC, "pockettts_spkenc"}, + { PROJECTOR_TYPE_POCKETTTS_GEN, "pockettts_gen"}, { PROJECTOR_TYPE_MUSE_GLIMMER, "muse-glimmer"}, }; diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 761aabf64e83..ad25c008e738 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -141,6 +141,20 @@ struct clip_hparams { int32_t rvq_num_quantizers = 0; std::vector rvq_codebook_size; // per-quantizer bin count (ragged, e.g. 1024/1024/256/128x17) + // threshold for the "out_eos_score" graph output + float gen_eos_threshold = 0.0f; + + // name of the weight variant, some pipelines tune themselves on it + std::string gen_model_variant; + + // pocket-tts + static constexpr int32_t pockettts_max_spk_seconds = 30; + int32_t seanet_n_stage = 0; + std::vector seanet_ratios; // encoder order (reversed compared to the config) + int32_t mimi_downsample = 0; // encoder frame rate / model frame rate + int32_t mimi_tfm_context = 0; // attention window of the mimi transformers, in frames + int32_t flow_n_step = 1; // lsd_decode steps + // qwen3tts code2wav int32_t wav_tfm_n_layer = 0; int32_t wav_tfm_n_embd = 0; @@ -402,6 +416,63 @@ struct qf_block { std::vector qf_proj_layers; }; +// pocket-tts SEANet stack, used in both directions: +// encoder = conv_in -> per stage (residual unit, strided conv) -> conv_out +// decoder = conv_in -> per stage (strided convtr, residual unit) -> conv_out +struct clip_seanet { + // one residual unit: ELU -> dilated conv -> ELU -> pointwise conv, added to the input + struct stage { + ggml_tensor * res_conv1_w = nullptr; + ggml_tensor * res_conv1_b = nullptr; + ggml_tensor * res_conv2_w = nullptr; + ggml_tensor * res_conv2_b = nullptr; + ggml_tensor * scale_conv_w = nullptr; // strided conv (encoder) or convtr (decoder) + ggml_tensor * scale_conv_b = nullptr; + }; + + ggml_tensor * conv_in_w = nullptr; + ggml_tensor * conv_in_b = nullptr; + ggml_tensor * conv_out_w = nullptr; + ggml_tensor * conv_out_b = nullptr; + std::vector stages; +}; + +// pocket-tts flow-matching decoder (SimpleMLPAdaLN) +struct clip_flow_net { + // AdaLN res block: in_ln -> modulate -> Linear -> SiLU -> Linear, gated residual + struct block { + ggml_tensor * norm_w = nullptr; + ggml_tensor * norm_b = nullptr; + ggml_tensor * up_w = nullptr; + ggml_tensor * up_b = nullptr; + ggml_tensor * down_w = nullptr; + ggml_tensor * down_b = nullptr; + ggml_tensor * ada_w = nullptr; // -> shift, scale, gate + ggml_tensor * ada_b = nullptr; + }; + + // timestep embedder: cos/sin(t * freqs) -> Linear -> SiLU -> Linear -> RMSNorm + struct time_embd { + ggml_tensor * freqs = nullptr; + ggml_tensor * up_w = nullptr; + ggml_tensor * up_b = nullptr; + ggml_tensor * down_w = nullptr; + ggml_tensor * down_b = nullptr; + ggml_tensor * norm = nullptr; // RMSNorm alpha + }; + + ggml_tensor * input_proj_w = nullptr; + ggml_tensor * input_proj_b = nullptr; + ggml_tensor * cond_embd_w = nullptr; + ggml_tensor * cond_embd_b = nullptr; + ggml_tensor * final_ada_w = nullptr; // -> shift, scale + ggml_tensor * final_ada_b = nullptr; + ggml_tensor * final_proj_w = nullptr; + ggml_tensor * final_proj_b = nullptr; + std::vector time; + std::vector blocks; +}; + // qwen3tts code2wav: RVQ codes -> raw PCM struct clip_code2wav { // "upsample" stage: one ConvNeXt block plus the causal ConvTranspose1d before it @@ -699,6 +770,24 @@ struct clip_model { // qwen3tts code2wav: RVQ codes -> raw PCM clip_code2wav c2w; + // pocket-tts: SEANet stack, shared by the encoder (speaker path) and the decoder (gen path) + clip_seanet seanet; + + // pocket-tts: voice latent -> backbone embd (speaker path) + ggml_tensor * spk_proj_w = nullptr; + ggml_tensor * downsample_w = nullptr; + + // pocket-tts: flow-matching decoder, backbone hidden state -> next latent + clip_flow_net flow; + ggml_tensor * gen_out_eos_w = nullptr; + ggml_tensor * gen_out_eos_b = nullptr; + ggml_tensor * gen_input_lin_w = nullptr; // latent -> backbone embd + ggml_tensor * gen_emb_mean = nullptr; + ggml_tensor * gen_emb_std = nullptr; + ggml_tensor * gen_quant_out_w = nullptr; // latent -> decoder dim + ggml_tensor * gen_upsample_w = nullptr; // depthwise convtr, frame rate -> encoder frame rate + std::vector gen_tfm_layers; // mimi decoder_transformer + // cogvlm ggml_tensor * mm_post_fc_norm_w = nullptr; ggml_tensor * mm_post_fc_norm_b = nullptr; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 1e53eddf8179..2fb2b5041dcd 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -174,6 +174,10 @@ struct clip_ctx { bool support_batch = false; + // for audio gen, reseeded only when the caller asks for another seed + std::mt19937 rng{std::random_device{}()}; + uint32_t rng_seed = UINT32_MAX; + clip_ctx(clip_context_params & ctx_params) { flash_attn_type = ctx_params.flash_attn_type; no_alloc = ctx_params.no_alloc; @@ -1059,6 +1063,25 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const { builder = std::make_unique(ctx, img); } break; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + { + builder = std::make_unique(ctx, img); + } break; + case PROJECTOR_TYPE_POCKETTTS_GEN: + { + const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE; + const int n_step = ctx->model.hparams.flow_n_step; + const int64_t n_latent = ctx->model.gen_input_lin_w->ne[0]; + GGML_ASSERT(n_step > 0); + GGML_ASSERT(n_latent > 0); + // "inp_feats" takes the caller's buffer as-is, the graph must consume all of it + if (params && params->feats) { + GGML_ASSERT(params->feats->size() % (size_t) n_latent == 0); + GGML_ASSERT(params->feats->size() >= (size_t) n_latent); + } + const int n_frames = params && params->feats ? (int) (params->feats->size() / n_latent) : 1; + builder = std::make_unique(ctx, img, gen_process, n_step, n_frames); + } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE; @@ -1282,6 +1305,7 @@ struct clip_model_loader { // these are unused, but still need to be set to avoid issues hparams.image_size = 0; hparams.patch_size = 1; + get_string(KEY_GEN_AUDIO_VARIANT, hparams.gen_model_variant, false); } else { GGML_ASSERT(false && "unknown modality"); @@ -1421,7 +1445,7 @@ struct clip_model_loader { } break; case PROJECTOR_TYPE_PARAKEET: { - get_u32(KEY_AUDIO_SUBSAMPLING_FACTOR, hparams.subsampling_factor); + get_u32(KEY_AUDIO_SUBSMPL_FACTOR, hparams.subsampling_factor); GGML_ASSERT(hparams.subsampling_factor == 8 && "subsampling_factor must match the conv strides in clip_graph_parakeet::build()"); get_u32(KEY_A_CONV_KERNEL_SIZE, hparams.audio_conv_kernel_size); @@ -1745,6 +1769,22 @@ struct clip_model_loader { // matches the reference decoder's sliding_window (speech_tokenizer/config.json) hparams.wav_tfm_swa = 72; } break; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + case PROJECTOR_TYPE_POCKETTTS_GEN: + { + // mimi front-end takes the raw waveform, no mel + hparams.audio_sample_rate = 24000; + // seanet ratios are [6,5,4] in the config, the encoder reverses them + hparams.seanet_ratios = { 4, 5, 6 }; + hparams.seanet_n_stage = (int32_t) hparams.seanet_ratios.size(); + hparams.mimi_downsample = 16; + // matches the reference transformer's "context" + hparams.mimi_tfm_context = 250; + hparams.rope_theta = 10000.0f; + // flow_lm defaults, see pocket_tts/default_parameters.py + hparams.flow_n_step = 1; + hparams.gen_eos_threshold = -4.0f; + } break; case PROJECTOR_TYPE_PADDLEOCR: { hparams.n_merge = 2; @@ -1947,7 +1987,9 @@ struct clip_model_loader { // GEMMA4UA is encoder-free: it uses n_mel_bins as a raw-waveform frame size (640) and has no FFT/filterbank, so the mel-range and FFT // checks below do not apply to it. - const bool fft_based = model.proj_type != PROJECTOR_TYPE_GEMMA4UA; + // pocket-tts is encoder-free in the same sense: mimi convolves the raw waveform + const bool fft_based = model.proj_type != PROJECTOR_TYPE_GEMMA4UA && + model.proj_type != PROJECTOR_TYPE_POCKETTTS_SPKENC; // Validate audio hparams loaded from GGUF metadata if (hparams.n_mel_bins <= 0 || (fft_based && hparams.n_mel_bins > 256)) { @@ -2020,6 +2062,31 @@ struct clip_model_loader { return cur; }; + // pocket-tts: the encoder and the decoder share the same layout, only the prefix differs + auto load_seanet = [&](clip_seanet & seanet, bool is_decoder) { + const char * conv_in = is_decoder ? TN_A_GEN_WAV_SEANET_CONV_IN : TN_A_SEANET_CONV_IN; + const char * conv_out = is_decoder ? TN_A_GEN_WAV_SEANET_CONV_OUT : TN_A_SEANET_CONV_OUT; + const char * res1 = is_decoder ? TN_A_GEN_WAV_SEANET_RES_CONV1 : TN_A_SEANET_RES_CONV1; + const char * res2 = is_decoder ? TN_A_GEN_WAV_SEANET_RES_CONV2 : TN_A_SEANET_RES_CONV2; + const char * scale = is_decoder ? TN_A_GEN_WAV_SEANET_SCALE_CONV : TN_A_SEANET_SCALE_CONV; + + seanet.conv_in_w = get_tensor(string_format(conv_in, "weight")); + seanet.conv_in_b = get_tensor(string_format(conv_in, "bias")); + seanet.conv_out_w = get_tensor(string_format(conv_out, "weight")); + seanet.conv_out_b = get_tensor(string_format(conv_out, "bias")); + + seanet.stages.resize(hparams.seanet_n_stage); + for (int i = 0; i < hparams.seanet_n_stage; i++) { + auto & stage = seanet.stages[i]; + stage.res_conv1_w = get_tensor(string_format(res1, i, "weight")); + stage.res_conv1_b = get_tensor(string_format(res1, i, "bias")); + stage.res_conv2_w = get_tensor(string_format(res2, i, "weight")); + stage.res_conv2_b = get_tensor(string_format(res2, i, "bias")); + stage.scale_conv_w = get_tensor(string_format(scale, i, "weight")); + stage.scale_conv_b = get_tensor(string_format(scale, i, "bias")); + } + }; + auto get_vector = [&](const std::string & name) { std::vector result; auto it = tensor_offset.find(name); @@ -2081,7 +2148,8 @@ struct clip_model_loader { const bool has_standard_layers = ( model.proj_type != PROJECTOR_TYPE_GEMMA3NV && - model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC); + model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC && + model.proj_type != PROJECTOR_TYPE_POCKETTTS_GEN); // layers const int n_layers_to_load = has_standard_layers ? hparams.n_layer : 0; @@ -2755,6 +2823,81 @@ struct clip_model_loader { model.mm_fc_w = get_tensor(string_format(TN_MM_AUDIO_FC, "weight")); model.mm_fc_b = get_tensor(string_format(TN_MM_AUDIO_FC, "bias")); } break; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + { + load_seanet(model.seanet, false); + model.downsample_w = get_tensor(string_format(TN_A_DOWNSAMPLE_CONV, "weight")); + model.spk_proj_w = get_tensor(string_format(TN_A_SPEAKER_PROJ, "weight")); + } break; + case PROJECTOR_TYPE_POCKETTTS_GEN: + { + auto & flow = model.flow; + flow.input_proj_w = get_tensor(string_format(TN_A_GEN_FLOW_INPUT_PROJ, "weight")); + flow.input_proj_b = get_tensor(string_format(TN_A_GEN_FLOW_INPUT_PROJ, "bias")); + flow.cond_embd_w = get_tensor(string_format(TN_A_GEN_FLOW_COND_EMBD, "weight")); + flow.cond_embd_b = get_tensor(string_format(TN_A_GEN_FLOW_COND_EMBD, "bias")); + flow.final_ada_w = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_ADA, "weight")); + flow.final_ada_b = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_ADA, "bias")); + flow.final_proj_w = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_PROJ, "weight")); + flow.final_proj_b = get_tensor(string_format(TN_A_GEN_FLOW_FINAL_PROJ, "bias")); + + flow.time.resize(2); + for (size_t i = 0; i < flow.time.size(); i++) { + auto & t = flow.time[i]; + t.freqs = get_tensor(string_format(TN_A_GEN_FLOW_TIME_FREQS, (int) i)); + t.up_w = get_tensor(string_format(TN_A_GEN_FLOW_TIME_UP, (int) i, "weight")); + t.up_b = get_tensor(string_format(TN_A_GEN_FLOW_TIME_UP, (int) i, "bias")); + t.down_w = get_tensor(string_format(TN_A_GEN_FLOW_TIME_DOWN, (int) i, "weight")); + t.down_b = get_tensor(string_format(TN_A_GEN_FLOW_TIME_DOWN, (int) i, "bias")); + t.norm = get_tensor(string_format(TN_A_GEN_FLOW_TIME_NORM, (int) i)); + } + + // one AdaLN block per flow depth, the count is only known from the tensors + for (int il = 0; ; il++) { + ggml_tensor * probe = get_tensor(string_format(TN_A_GEN_FLOW_BLK_NORM, il, "weight"), false); + if (probe == nullptr) { + break; + } + clip_flow_net::block blk; + blk.norm_w = probe; + blk.norm_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_NORM, il, "bias")); + blk.up_w = get_tensor(string_format(TN_A_GEN_FLOW_BLK_UP, il, "weight")); + blk.up_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_UP, il, "bias")); + blk.down_w = get_tensor(string_format(TN_A_GEN_FLOW_BLK_DOWN, il, "weight")); + blk.down_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_DOWN, il, "bias")); + blk.ada_w = get_tensor(string_format(TN_A_GEN_FLOW_BLK_ADA, il, "weight")); + blk.ada_b = get_tensor(string_format(TN_A_GEN_FLOW_BLK_ADA, il, "bias")); + flow.blocks.push_back(blk); + } + + model.gen_out_eos_w = get_tensor(string_format(TN_A_GEN_OUT_EOS, "weight")); + model.gen_out_eos_b = get_tensor(string_format(TN_A_GEN_OUT_EOS, "bias")); + model.gen_input_lin_w = get_tensor(string_format(TN_A_GEN_INPUT_LINEAR, "weight")); + model.gen_emb_mean = get_tensor(TN_A_GEN_EMB_MEAN); + model.gen_emb_std = get_tensor(TN_A_GEN_EMB_STD); + + // mimi decoder + model.gen_quant_out_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_OUT, "weight")); + model.gen_upsample_w = get_tensor(string_format(TN_A_GEN_WAV_UPSAMPLE, "weight")); + load_seanet(model.seanet, true); + model.gen_tfm_layers.resize(hparams.n_layer); + for (int il = 0; il < hparams.n_layer; il++) { + auto & layer = model.gen_tfm_layers[il]; + const char * p = "a.gen.wav.tfm"; + layer.ln_1_w = get_tensor(string_format(TN_LN_1, p, il, "weight")); + layer.ln_1_b = get_tensor(string_format(TN_LN_1, p, il, "bias")); + layer.q_w = get_tensor(string_format(TN_ATTN_Q, p, il, "weight")); + layer.k_w = get_tensor(string_format(TN_ATTN_K, p, il, "weight")); + layer.v_w = get_tensor(string_format(TN_ATTN_V, p, il, "weight")); + layer.o_w = get_tensor(string_format(TN_ATTN_OUTPUT, p, il, "weight")); + layer.ls_1_w = get_tensor(string_format(TN_LS_1, p, il, "weight")); + layer.ln_2_w = get_tensor(string_format(TN_LN_2, p, il, "weight")); + layer.ln_2_b = get_tensor(string_format(TN_LN_2, p, il, "bias")); + layer.ff_up_w = get_tensor(string_format(TN_FFN_UP, p, il, "weight")); + layer.ff_down_w = get_tensor(string_format(TN_FFN_DOWN, p, il, "weight")); + layer.ls_2_w = get_tensor(string_format(TN_LS_2, p, il, "weight")); + } + } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { // code_predictor @@ -4060,6 +4203,17 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { // one hidden-state vector fed back to the talker per call n_patches = 1; } break; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + { + // one conditioning row per 12.5Hz frame + const int hop = ctx->model.hparams.mimi_downsample * 120; + n_patches = img->nx() / hop; + } break; + case PROJECTOR_TYPE_POCKETTTS_GEN: + { + // one latent per call for GEN_CODE, GEN_WAV sizes its input from the caller + n_patches = 1; + } break; case PROJECTOR_TYPE_GRANITE4_VISION: { // Per-tile output token count: each projector block outputs @@ -4101,6 +4255,15 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32 return clip_encode(ctx, ¶ms); } +// persisted state slots of the gen-audio decoder, per pipeline +static std::vector list_gen_state_slots(const clip_hparams & hparams, const clip_model & model) { + switch (model.proj_type) { + case PROJECTOR_TYPE_QWEN3TTS_GEN: return list_c2w_state_slots(hparams, model); + case PROJECTOR_TYPE_POCKETTTS_GEN: return list_pockettts_state_slots(hparams, model); + default: return {}; + } +} + bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { const clip_image_f32_batch & imgs = *params->imgs; int n_batch_cur = imgs.entries.size(); @@ -4116,6 +4279,11 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { clip_model_loader::warmup(*ctx, *params->imgs); } + if (params->seed != ctx->rng_seed) { + ctx->rng_seed = params->seed; + ctx->rng.seed(params->seed == UINT32_MAX ? std::random_device{}() : params->seed); + } + // build the inference graph ggml_backend_sched_reset(ctx->sched.get()); ggml_cgraph * gf = clip_get_graph_builder(ctx, imgs, params)->build(); @@ -4160,6 +4328,50 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { ggml_backend_tensor_set(cur, values.data(), 0, ggml_nbytes(cur)); }; + // upload the decoder state from the previous call, or zero-fill on a cold start + auto set_gen_state_in = [&]() { + size_t offset = 0; + for (const auto & slot : list_gen_state_slots(hparams, model)) { + ggml_tensor * t = get_inp_tensor(("state_in_" + slot.name).c_str()); + const size_t nb = ggml_nbytes(t); + if (params->state_in && params->state_in->size() >= offset + nb) { + ggml_backend_tensor_set(t, params->state_in->data() + offset, 0, nb); + } else { + std::vector zeros(nb, 0); + ggml_backend_tensor_set(t, zeros.data(), 0, nb); + } + offset += nb; + } + }; + + // rope positions and attention mask of the mimi transformers (pocket-tts). + // the mask is causal with a sliding window, see _build_attention_mask() in the reference + auto set_pockettts_tfm_inputs = [&]() { + const int64_t n_pos = ggml_nelements(get_inp_tensor("inp_pos")); + GGML_ASSERT(n_pos > 0); + std::vector positions((size_t) n_pos); + for (int64_t i = 0; i < n_pos; i++) { + positions[(size_t) i] = (int32_t) i; + } + set_input_i32("inp_pos", positions); + + // the preprocessor truncates the waveform to keep this mask bounded + const int64_t max_pos = (int64_t) clip_hparams::pockettts_max_spk_seconds * hparams.audio_sample_rate / 120; + GGML_ASSERT(n_pos <= max_pos && "pocket-tts speaker reference too long for a dense mask"); + + const int64_t context = hparams.mimi_tfm_context; + std::vector mask((size_t) n_pos * n_pos, -INFINITY); + for (int64_t q = 0; q < n_pos; q++) { + for (int64_t k = 0; k < n_pos; k++) { + const int64_t delta = q - k; + if (delta >= 0 && delta < context) { + mask[(size_t) q * n_pos + k] = 0.0f; + } + } + } + set_input_f32("kq_mask", mask); + }; + // set input pixel values if (!imgs.is_audio) { size_t nelem = 0; @@ -4203,8 +4415,8 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } set_input_f32("inp_raw", inp_raw); - } else if (!(ctx->proj_type() == PROJECTOR_TYPE_QWEN3TTS_GEN && params->gen_process == CLIP_GEN_PROCESS_GEN_WAV)) { - // audio input, code2wav is not here: its only input is "inp_codes", set in the switch below + } else if (params->gen_process != CLIP_GEN_PROCESS_GEN_WAV) { + // audio input. GEN_WAV is not here: it takes codes or feats, set in the switch below GGML_ASSERT(imgs.entries.size() == 1); const auto & mel_inp = imgs.entries[0]; @@ -4737,6 +4949,30 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } set_input_i32("patches", patches); } break; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + { + set_pockettts_tfm_inputs(); + } break; + case PROJECTOR_TYPE_POCKETTTS_GEN: + { + if (params->gen_process == CLIP_GEN_PROCESS_GEN_WAV) { + GGML_ASSERT(params->feats != nullptr); + set_input_f32("inp_feats", *params->feats); + // positions and mask are derived in-graph from the persisted counter + set_gen_state_in(); + } else { + // flow matching starts from gaussian noise, std = sqrt(temp) + ggml_tensor * t = get_inp_tensor("inp_noise"); + // Config.default_temperature, for a caller that does not set one + const float temp = params->temp > 0.0f ? params->temp : 0.7f; + std::normal_distribution dist(0.0f, std::sqrt(temp)); + std::vector noise(ggml_nelements(t)); + for (auto & v : noise) { + v = dist(ctx->rng); + } + set_input_f32("inp_noise", noise); + } + } break; case PROJECTOR_TYPE_GEMMA4V: case PROJECTOR_TYPE_GEMMA4UV: { @@ -4861,20 +5097,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } } set_input_i32("inp_codes", codes); - - // upload the state from the previous call, or zero-fill on a cold start - size_t offset = 0; - for (const auto & slot : list_c2w_state_slots(hparams, model)) { - ggml_tensor * t = get_inp_tensor(("state_in_" + slot.name).c_str()); - const size_t nb = ggml_nbytes(t); - if (params->state_in && params->state_in->size() >= offset + nb) { - ggml_backend_tensor_set(t, params->state_in->data() + offset, 0, nb); - } else { - std::vector zeros(nb, 0); - ggml_backend_tensor_set(t, zeros.data(), 0, nb); - } - offset += nb; - } + set_gen_state_in(); } else { // code0 indexes gen_code_out_embd_w via ggml_get_rows; bound it const int64_t vocab0 = model.gen_code_out_embd_w->ne[1]; @@ -4886,11 +5109,10 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { set_input_i32("inp_code0", code0); // one uniform(0,1) draw per codebook, used by do_sampling() - static std::mt19937 rng{ std::random_device{}() }; std::uniform_real_distribution dist(0.0f, 1.0f); const int64_t n_acoustic = model.gen_code_head_w->ne[2]; for (int64_t g = 0; g < n_acoustic; g++) { - std::vector r = { dist(rng) }; + std::vector r = { dist(ctx->rng) }; set_input_f32(("inp_rand_" + std::to_string(g)).c_str(), r); } } @@ -5343,14 +5565,31 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { // for audio gen models // + // optional outputs: a pipeline yields codes or feats, and not all have an eos head if (params->out_codes != nullptr) { ggml_tensor * codes = ggml_graph_get_tensor(gf, "out_codes"); - if (codes == nullptr) { - GGML_ABORT("out_codes requested but graph has no \"out_codes\" tensor"); + if (codes != nullptr) { + auto & out_codes = *params->out_codes; + out_codes.resize(ggml_nelements(codes)); + ggml_backend_tensor_get(codes, out_codes.data(), 0, ggml_nbytes(codes)); + } + } + if (params->out_feats != nullptr) { + ggml_tensor * feats = ggml_graph_get_tensor(gf, "out_feats"); + if (feats != nullptr) { + auto & out_feats = *params->out_feats; + out_feats.resize(ggml_nelements(feats)); + ggml_backend_tensor_get(feats, out_feats.data(), 0, ggml_nbytes(feats)); + } + } + if (params->out_is_eos != nullptr) { + ggml_tensor * eos = ggml_graph_get_tensor(gf, "out_eos_score"); + if (eos != nullptr) { + GGML_ASSERT(ggml_nelements(eos) == 1); + float score = 0.0f; + ggml_backend_tensor_get(eos, &score, 0, sizeof(float)); + *params->out_is_eos = score > hparams.gen_eos_threshold; } - auto & out_codes = *params->out_codes; - out_codes.resize(ggml_nelements(codes)); - ggml_backend_tensor_get(codes, out_codes.data(), 0, ggml_nbytes(codes)); } if (params->out_audio != nullptr) { ggml_tensor * audio = ggml_graph_get_tensor(gf, "out_audio"); @@ -5362,9 +5601,9 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { ggml_backend_tensor_get(audio, out_audio.data(), 0, ggml_nbytes(audio)); // drop the tail audio that comes from the code-0 rear padding - const int64_t n_codes = model.gen_code_head_w->ne[2] + 1; + const int64_t n_codes = params->codes ? model.gen_code_head_w->ne[2] + 1 : 0; const int64_t n_frames_w = hparams.wav_tfm_swa; - const int64_t n_frames = (int64_t) params->codes->size() / n_codes; + const int64_t n_frames = params->codes ? (int64_t) params->codes->size() / n_codes : n_frames_w; if (n_frames < n_frames_w) { const size_t hop = out_audio.size() / n_frames_w; out_audio.resize((size_t) n_frames * hop); @@ -5373,12 +5612,12 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { if (params->state_out != nullptr) { auto & state_out = *params->state_out; size_t total = 0; - for (const auto & slot : list_c2w_state_slots(hparams, model)) { + for (const auto & slot : list_gen_state_slots(hparams, model)) { total += (size_t) (slot.ne0 * slot.ne1) * sizeof(float); } state_out.resize(total); size_t offset = 0; - for (const auto & slot : list_c2w_state_slots(hparams, model)) { + for (const auto & slot : list_gen_state_slots(hparams, model)) { ggml_tensor * t = ggml_graph_get_tensor(gf, ("state_out_" + slot.name).c_str()); if (t == nullptr) { GGML_ABORT("state_out requested but graph has no \"state_out_%s\" tensor", slot.name.c_str()); @@ -5526,6 +5765,10 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { return ctx->model.mm_fc_w->ne[2]; case PROJECTOR_TYPE_QWEN3TTS_GEN: return ctx->model.gen_code_out_embd_w->ne[0]; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + return ctx->model.spk_proj_w->ne[1]; + case PROJECTOR_TYPE_POCKETTTS_GEN: + return ctx->model.gen_input_lin_w->ne[1]; case PROJECTOR_TYPE_PARAKEET: return ctx->model.mm_1_w->ne[1]; default: diff --git a/tools/mtmd/clip.h b/tools/mtmd/clip.h index 7f706d976eb3..a5b713775234 100644 --- a/tools/mtmd/clip.h +++ b/tools/mtmd/clip.h @@ -104,9 +104,14 @@ struct clip_encode_params { int32_t top_k = 50; float top_p = 1.0f; std::vector * out_codes = nullptr; // this frame's 16 sampled codes + std::vector * out_feats = nullptr; // continuous counterpart of out_codes + uint32_t seed = UINT32_MAX; // UINT32_MAX for random + float temp = 0.0f; // sampling temperature, noise scale for flow-matching decoders + bool * out_is_eos = nullptr; // GEN_WAV const std::vector * codes = nullptr; // this frame's 16 RVQ codes + const std::vector * feats = nullptr; // continuous counterpart of codes std::vector * out_audio = nullptr; // decoded PCM samples, F32 const std::vector * state_in = nullptr; // state from previous call, null or wrong size means cold start std::vector * state_out = nullptr; // state for the next call diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index 519c0d019add..ed8c1ea51870 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -318,6 +318,59 @@ struct clip_graph_qwen3tts_gen : clip_graph { }; }; +// +// pocket-tts: SEANet convolution stack, shared by the voice encoder and the mimi decoder. +// stateless unless state_in is populated: convs then pad instead of carrying left-context. +// +struct clip_graph_pockettts_seanet : clip_graph { + clip_graph_pockettts_seanet(const clip_graph & parent) : clip_graph(parent) {} + ggml_cgraph * build() override { GGML_ABORT("call encode()/decode() instead"); } + + // per-call streaming state, keyed by slot name (see list_pockettts_state_slots) + std::map state_in; + mutable std::vector> state_out; + + ggml_tensor * conv1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, int dilation, + bool pad_replicate = false, const std::string & state_name = "") const; + ggml_tensor * conv_transpose1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, + const std::string & state_name = "") const; + ggml_tensor * res_unit(ggml_tensor * x, const clip_seanet::stage & stage, int dilation, + const std::string & state_prefix = "") const; + + // x: [T, C] -> [T / hop, dim] + ggml_tensor * encode(ggml_tensor * x) const; + // x: [T, dim] -> [T * hop, 1], streams when state_in is populated + ggml_tensor * decode(ggml_tensor * x) const; +}; + +// mimi encoder + speaker_proj: reference waveform -> voice conditioning rows +struct clip_graph_pockettts_spkenc : clip_graph { + clip_graph_pockettts_spkenc(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} + ggml_cgraph * build() override; + + ggml_tensor * tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, ggml_tensor * inp_pos, ggml_tensor * kq_mask, int il) const; +}; + +// +// pocket-tts generation: +// GEN_CODE = flow-matching decoder + end-of-speech head, one latent per call +// GEN_WAV = mimi decoder, a window of latents -> PCM +// +struct clip_graph_pockettts_gen : clip_graph { + clip_graph_pockettts_gen(clip_ctx * ctx, const clip_image_f32 & img, clip_gen_process_type gen_process, int n_step, int n_frames) + : clip_graph(ctx, img), gen_process(gen_process), n_step(n_step), n_frames(n_frames) {} + ggml_cgraph * build() override; + + clip_gen_process_type gen_process; + int n_step; // lsd_decode steps, fixed at graph-build time + int n_frames; // GEN_WAV only: number of latents to decode + + // AdaLN modulation: x * (1 + scale) + shift + ggml_tensor * modulate(ggml_tensor * x, ggml_tensor * shift, ggml_tensor * scale) const; + ggml_tensor * time_embed(const clip_flow_net::time_embd & te, float t) const; + ggml_tensor * flow_forward(ggml_tensor * cond, ggml_tensor * x, float s, float t) const; +}; + // one persisted state buffer used by code2wav, see qwen3tts-gen.cpp struct c2w_state_slot { std::string name; @@ -326,6 +379,9 @@ struct c2w_state_slot { }; std::vector list_c2w_state_slots(const clip_hparams & hparams, const clip_model & model); +// same, for the streaming mimi decoder (pocket-tts GEN_WAV) +std::vector list_pockettts_state_slots(const clip_hparams & hparams, const clip_model & model); + struct clip_graph_kimik25 : clip_graph { clip_graph_kimik25(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; diff --git a/tools/mtmd/models/pockettts-gen.cpp b/tools/mtmd/models/pockettts-gen.cpp new file mode 100644 index 000000000000..3fd613e5f7fa --- /dev/null +++ b/tools/mtmd/models/pockettts-gen.cpp @@ -0,0 +1,291 @@ +#include "models.h" + +#include + +// pocket-tts generation stages +// +// GEN_CODE: backbone hidden state -> next 32-d latent (flow matching) + end-of-speech score +// GEN_WAV : a window of latents -> PCM, through the mimi decoder +// +// there is no codebook anywhere, "codes" in the mtmd API are continuous features here + +ggml_tensor * clip_graph_pockettts_gen::modulate(ggml_tensor * x, ggml_tensor * shift, ggml_tensor * scale) const { + ggml_tensor * cur = ggml_mul(ctx0, x, ggml_scale_bias(ctx0, scale, 1.0f, 1.0f)); + return ggml_add(ctx0, cur, shift); +} + +// see TimestepEmbedder in the reference +ggml_tensor * clip_graph_pockettts_gen::time_embed(const clip_flow_net::time_embd & te, float t) const { + // t is a graph-build constant, so the cos/sin table can be folded into a scaled copy + ggml_tensor * args = ggml_scale(ctx0, te.freqs, t); + ggml_tensor * emb = ggml_concat(ctx0, ggml_cos(ctx0, args), ggml_sin(ctx0, args), 0); + + ggml_tensor * cur = build_mm(te.up_w, emb); + cur = ggml_add(ctx0, cur, te.up_b); + cur = ggml_silu(ctx0, cur); + cur = build_mm(te.down_w, cur); + cur = ggml_add(ctx0, cur, te.down_b); + + // this "RMSNorm" divides by the unbiased variance, not the mean square + // it also rescales the input, not the centered value, see _rms_norm() in mlp.py + { + const int64_t n = cur->ne[0]; + ggml_tensor * mean = ggml_mean(ctx0, cur); + ggml_tensor * dev = ggml_sub(ctx0, cur, mean); + ggml_tensor * var = ggml_mean(ctx0, ggml_sqr(ctx0, dev)); + var = ggml_scale_bias(ctx0, var, (float) n / (float) (n - 1), 1e-5f); + cur = ggml_div(ctx0, cur, ggml_sqrt(ctx0, var)); + cur = ggml_mul(ctx0, cur, te.norm); + } + + return cur; +} + +// one velocity evaluation: v(cond, s, t, x) +ggml_tensor * clip_graph_pockettts_gen::flow_forward(ggml_tensor * cond, ggml_tensor * x, float s, float t) const { + const auto & flow = model.flow; + + ggml_tensor * cur = build_mm(flow.input_proj_w, x); + cur = ggml_add(ctx0, cur, flow.input_proj_b); + + // the two time conditions are averaged, then added to the projected backbone state + ggml_tensor * ts = ggml_add(ctx0, time_embed(flow.time[0], s), time_embed(flow.time[1], t)); + ts = ggml_scale(ctx0, ts, 1.0f / (float) flow.time.size()); + + ggml_tensor * c = build_mm(flow.cond_embd_w, cond); + c = ggml_add(ctx0, c, flow.cond_embd_b); + + ggml_tensor * y = ggml_add(ctx0, ts, c); + cb(y, "flow_cond", -1); + + const int64_t n_ch = flow.blocks.empty() ? 0 : flow.blocks[0].norm_w->ne[0]; + + for (size_t il = 0; il < flow.blocks.size(); il++) { + const auto & blk = flow.blocks[il]; + + ggml_tensor * mod = build_mm(blk.ada_w, ggml_silu(ctx0, y)); + mod = ggml_add(ctx0, mod, blk.ada_b); + + ggml_tensor * shift = ggml_view_1d(ctx0, mod, n_ch, 0); + ggml_tensor * scale = ggml_view_1d(ctx0, mod, n_ch, (size_t) n_ch * mod->nb[0]); + ggml_tensor * gate = ggml_view_1d(ctx0, mod, n_ch, (size_t) 2 * n_ch * mod->nb[0]); + + ggml_tensor * h = build_norm(cur, blk.norm_w, blk.norm_b, NORM_TYPE_NORMAL, 1e-6f, (int) il); + h = modulate(h, shift, scale); + h = build_mm(blk.up_w, h); + h = ggml_add(ctx0, h, blk.up_b); + h = ggml_silu(ctx0, h); + h = build_mm(blk.down_w, h); + h = ggml_add(ctx0, h, blk.down_b); + + cur = ggml_add(ctx0, cur, ggml_mul(ctx0, gate, h)); + cb(cur, "flow_blk", (int) il); + } + + // final layer: the norm has no weights, only the AdaLN modulation + ggml_tensor * mod = build_mm(flow.final_ada_w, ggml_silu(ctx0, y)); + mod = ggml_add(ctx0, mod, flow.final_ada_b); + + ggml_tensor * shift = ggml_view_1d(ctx0, mod, n_ch, 0); + ggml_tensor * scale = ggml_view_1d(ctx0, mod, n_ch, (size_t) n_ch * mod->nb[0]); + + cur = build_norm(cur, nullptr, nullptr, NORM_TYPE_NORMAL, 1e-6f, -1); + cur = modulate(cur, shift, scale); + cur = build_mm(flow.final_proj_w, cur); + cur = ggml_add(ctx0, cur, flow.final_proj_b); + + return cur; +} + +// state carried between GEN_WAV calls: rope offset, per-layer KV window, conv left context +// and the transposed-conv overlap tails +std::vector list_pockettts_state_slots(const clip_hparams & hparams, const clip_model & model) { + std::vector slots; + if (model.gen_upsample_w == nullptr) { + return slots; // not a pocket-tts decoder + } + const auto & seanet = model.seanet; + + // the slots below are sized from these + GGML_ASSERT(!model.gen_tfm_layers.empty()); + GGML_ASSERT((int) seanet.stages.size() >= hparams.seanet_n_stage); + GGML_ASSERT((int) hparams.seanet_ratios.size() >= hparams.seanet_n_stage); + GGML_ASSERT(hparams.mimi_tfm_context > 1 && hparams.mimi_downsample > 0); + + slots.push_back({"tfm_pos", 1, 1}); + + const int64_t n_embd_a = model.gen_tfm_layers[0].q_w->ne[1]; + const int64_t prefix = hparams.mimi_tfm_context - 1; + for (size_t il = 0; il < model.gen_tfm_layers.size(); il++) { + slots.push_back({"tfm_k_" + std::to_string(il), n_embd_a, prefix}); + slots.push_back({"tfm_v_" + std::to_string(il), n_embd_a, prefix}); + } + + // upsample is depthwise, its output channel count is the input one + slots.push_back({"up", model.gen_upsample_w->ne[0] - hparams.mimi_downsample, model.gen_upsample_w->ne[2]}); + + slots.push_back({"dec_in", seanet.conv_in_w->ne[0] - 1, seanet.conv_in_w->ne[1]}); + for (int i = 0; i < hparams.seanet_n_stage; i++) { + const auto & stage = seanet.stages[i]; + const int stride = hparams.seanet_ratios[hparams.seanet_n_stage - 1 - i]; + slots.push_back({"dec_up_" + std::to_string(i), stage.scale_conv_w->ne[0] - stride, stage.scale_conv_w->ne[1]}); + slots.push_back({"dec_res_" + std::to_string(i), stage.res_conv1_w->ne[0] - 1, stage.res_conv1_w->ne[1]}); + } + slots.push_back({"dec_out", seanet.conv_out_w->ne[0] - 1, seanet.conv_out_w->ne[1]}); + + return slots; +} + +ggml_cgraph * clip_graph_pockettts_gen::build() { + if (gen_process == CLIP_GEN_PROCESS_GEN_CODE) { + // the backbone hidden state arrives as the single batch entry + ggml_tensor * h_state = build_inp_raw(1); + h_state = ggml_reshape_2d(ctx0, h_state, n_mmproj_embd, 1); + + // end-of-speech probe, thresholded on the host side + ggml_tensor * eos = build_mm(model.gen_out_eos_w, h_state); + eos = ggml_add(ctx0, eos, model.gen_out_eos_b); + ggml_set_name(eos, "out_eos_score"); + ggml_set_output(eos); + ggml_build_forward_expand(gf, eos); + + const int64_t n_latent = model.gen_input_lin_w->ne[0]; + + ggml_tensor * noise = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_latent, 1); + ggml_set_name(noise, "inp_noise"); + ggml_set_input(noise); + + // lsd_decode: integrate the velocity field from the noise sample + ggml_tensor * cur = noise; + for (int i = 0; i < n_step; i++) { + const float s = (float) i / (float) n_step; + const float t = (float) (i + 1) / (float) n_step; + ggml_tensor * v = flow_forward(h_state, cur, s, t); + cur = ggml_add(ctx0, cur, ggml_scale(ctx0, v, 1.0f / (float) n_step)); + } + cb(cur, "flow_latent", -1); + + ggml_set_name(cur, "out_feats"); + ggml_set_output(cur); + ggml_build_forward_expand(gf, cur); + + // the same latent, projected into the backbone's input space for the next step + ggml_tensor * embd = build_mm(model.gen_input_lin_w, cur); + cb(embd, "gen_embd", -1); + ggml_build_forward_expand(gf, embd); + + return gf; + } + + // GEN_WAV: [32, n_frames] latents -> PCM + ggml_tensor * feats = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, + model.gen_input_lin_w->ne[0], n_frames); + ggml_set_name(feats, "inp_feats"); + ggml_set_input(feats); + + // denormalize, then the DummyQuantizer up-projection + ggml_tensor * cur = ggml_add(ctx0, ggml_mul(ctx0, feats, model.gen_emb_std), model.gen_emb_mean); + cur = build_mm(model.gen_quant_out_w, cur); + cb(cur, "quant_out", -1); + + clip_graph_pockettts_seanet seanet(*this); + for (const auto & slot : list_pockettts_state_slots(hparams, model)) { + ggml_tensor * t = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, slot.ne0, slot.ne1); + ggml_set_name(t, ("state_in_" + slot.name).c_str()); + ggml_set_input(t); + seanet.state_in[slot.name] = t; + } + + // model frame rate -> encoder frame rate, depthwise transposed conv + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + cur = seanet.conv_transpose1d(cur, model.gen_upsample_w, nullptr, hparams.mimi_downsample, "up"); + cb(cur, "mimi_upsample", -1); + + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + + // positions continue across calls, the counter lives in the state + const int64_t n_pos = cur->ne[1]; + const int64_t prefix = hparams.mimi_tfm_context - 1; + const int64_t n_kv = prefix + n_pos; + + ggml_tensor * base = ggml_reshape_1d(ctx0, seanet.state_in.at("tfm_pos"), 1); + ggml_tensor * inp_pos = ggml_cast(ctx0, ggml_add(ctx0, ggml_arange(ctx0, 0.0f, (float) n_pos, 1.0f), base), + GGML_TYPE_I32); + seanet.state_out.push_back({"tfm_pos", ggml_scale_bias(ctx0, seanet.state_in.at("tfm_pos"), 1.0f, (float) n_pos)}); + + // banded causal mask over [cached prefix | this chunk] + // the last factor masks out cache rows that hold no real frame yet + ggml_tensor * pos_k = ggml_reshape_2d(ctx0, ggml_arange(ctx0, 0.0f, (float) n_kv, 1.0f), n_kv, 1); + ggml_tensor * pos_q = ggml_reshape_2d(ctx0, ggml_arange(ctx0, (float) prefix, (float) (prefix + n_pos), 1.0f), 1, n_pos); + ggml_tensor * diff = ggml_sub(ctx0, ggml_repeat_4d(ctx0, pos_q, n_kv, n_pos, 1, 1), pos_k); + + ggml_tensor * keep = ggml_mul(ctx0, + ggml_step(ctx0, ggml_scale_bias(ctx0, diff, 1.0f, 0.5f)), // delta >= 0 + ggml_step(ctx0, ggml_scale_bias(ctx0, diff, -1.0f, (float) hparams.mimi_tfm_context - 0.5f))); // delta < context + keep = ggml_mul(ctx0, keep, + ggml_step(ctx0, ggml_scale_bias(ctx0, ggml_add(ctx0, pos_k, base), 1.0f, 0.5f - (float) prefix))); + ggml_tensor * kq_mask = ggml_reshape_4d(ctx0, ggml_log(ctx0, keep), n_kv, n_pos, 1, 1); + + for (int il = 0; il < n_layer; il++) { + const auto & layer = model.gen_tfm_layers[il]; + ggml_tensor * inp = cur; + + cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, eps, il); + + ggml_tensor * Qcur = build_mm(layer.q_w, cur); + ggml_tensor * Kcur = build_mm(layer.k_w, cur); + ggml_tensor * Vcur = build_mm(layer.v_w, cur); + + Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos); + Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos); + + Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0, + hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0, + hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + + // prepend the cached window, then keep this chunk's tail for the next call + const std::string k_name = "tfm_k_" + std::to_string(il); + const std::string v_name = "tfm_v_" + std::to_string(il); + ggml_tensor * k_full = ggml_concat(ctx0, seanet.state_in.at(k_name), + ggml_reshape_2d(ctx0, Kcur, d_head * n_head, n_pos), 1); + ggml_tensor * v_full = ggml_concat(ctx0, seanet.state_in.at(v_name), Vcur, 1); + seanet.state_out.push_back({k_name, ggml_cont(ctx0, ggml_view_2d(ctx0, k_full, k_full->ne[0], prefix, + k_full->nb[1], (size_t) n_pos * k_full->nb[1]))}); + seanet.state_out.push_back({v_name, ggml_cont(ctx0, ggml_view_2d(ctx0, v_full, v_full->ne[0], prefix, + v_full->nb[1], (size_t) n_pos * v_full->nb[1]))}); + + ggml_tensor * q_cur = ggml_reshape_4d(ctx0, Qcur, d_head, n_head, n_pos, 1); + ggml_tensor * k_cur = ggml_reshape_4d(ctx0, k_full, d_head, n_head, n_kv, 1); + ggml_tensor * v_cur = ggml_reshape_4d(ctx0, v_full, d_head, n_head, n_kv, 1); + + cur = build_attn(layer.o_w, nullptr, q_cur, k_cur, v_cur, kq_mask, kq_scale, il); + cur = ggml_mul(ctx0, cur, layer.ls_1_w); + cur = ggml_add(ctx0, cur, inp); + + inp = cur; + cur = build_norm(cur, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, eps, il); + cur = build_ffn(cur, layer.ff_up_w, nullptr, nullptr, nullptr, layer.ff_down_w, nullptr, FFN_GELU, il); + cur = ggml_mul(ctx0, cur, layer.ls_2_w); + cur = ggml_add(ctx0, cur, inp); + } + cb(cur, "mimi_dec_tfm", -1); + + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + cur = seanet.decode(cur); + + for (const auto & s : seanet.state_out) { + ggml_set_name(s.second, ("state_out_" + s.first).c_str()); + ggml_set_output(s.second); + ggml_build_forward_expand(gf, s.second); + } + + // [n_samples, 1] -> [n_samples], clamped like the reference output + cur = ggml_reshape_1d(ctx0, cur, cur->ne[0]); + cur = ggml_clamp(ctx0, cur, -1.0f, 1.0f); + ggml_set_name(cur, "out_audio"); + ggml_set_output(cur); + ggml_build_forward_expand(gf, cur); + + return gf; +} diff --git a/tools/mtmd/models/pockettts-seanet.cpp b/tools/mtmd/models/pockettts-seanet.cpp new file mode 100644 index 000000000000..c47207f569bb --- /dev/null +++ b/tools/mtmd/models/pockettts-seanet.cpp @@ -0,0 +1,162 @@ +#include "models.h" + +// SEANet convolution stack of the mimi codec, see pocket_tts/modules/seanet.py +// +// tensors are T-first here: [T, C] +// the convs are causal: left context comes from a state slot, or from padding on a cold start + +static int64_t div_ceil(int64_t a, int64_t b) { + return a / b + (a % b ? 1 : 0); +} + +// x: [T, IC], w: [K, IC, OC] -> [T / stride, OC] +// the convs are causal, so the whole K - stride padding goes on the left +ggml_tensor * clip_graph_pockettts_seanet::conv1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, int dilation, + bool pad_replicate, const std::string & state_name) const { + const int64_t k_size = (w->ne[0] - 1) * dilation + 1; + const int64_t p_total = k_size - stride; + + // trailing padding so the last frame is not dropped, see pad_for_conv1d() in conv.py + const int64_t n_frames = div_ceil(x->ne[0] - k_size + p_total, stride); + const int64_t ideal_len = n_frames * stride + k_size - p_total; + const int64_t p_extra = ideal_len - x->ne[0]; + + if (!state_name.empty() && p_total > 0) { + // streaming: the left context is the tail of the previous call + ggml_tensor * left = state_in.at(state_name); // [p_total, IC] + x = ggml_concat(ctx0, left, x, 0); + state_out.push_back({state_name, + ggml_cont(ctx0, ggml_view_2d(ctx0, x, p_total, x->ne[1], x->nb[1], + (size_t) (x->ne[0] - p_total) * x->nb[0]))}); + } else if (pad_replicate && p_total > 0) { + // the resamplers repeat the first frame instead of zero-padding + ggml_tensor * first = ggml_view_2d(ctx0, x, 1, x->ne[1], x->nb[1], 0); + ggml_tensor * left = ggml_repeat_4d(ctx0, first, p_total, x->ne[1], 1, 1); + x = ggml_concat(ctx0, left, x, 0); + x = ggml_pad_ext(ctx0, x, 0, p_extra, 0, 0, 0, 0, 0, 0); + } else { + x = ggml_pad_ext(ctx0, x, p_total, p_extra, 0, 0, 0, 0, 0, 0); + } + + ggml_tensor * y = ggml_conv_1d(ctx0, w, x, stride, 0, dilation); + y = ggml_reshape_2d(ctx0, y, y->ne[0], y->ne[1]); + if (b) { + y = ggml_add(ctx0, y, ggml_reshape_2d(ctx0, b, 1, b->ne[0])); + } + return y; +} + +// x: [T, IC], w: [K, OC/groups, IC] -> [T * stride, OC] +// the K - stride overlap tail belongs to the next call: added to its head when streaming, else dropped +ggml_tensor * clip_graph_pockettts_seanet::conv_transpose1d(ggml_tensor * x, ggml_tensor * w, ggml_tensor * b, int stride, + const std::string & state_name) const { + const int64_t K = w->ne[0]; + const int64_t T = x->ne[0]; + const int64_t p_total = K - stride; + const bool depthwise = w->ne[1] == 1 && w->ne[2] > 1; + const int64_t OC = depthwise ? w->ne[2] : w->ne[1]; + const int64_t emit_len = T * stride; + + // one column per input step, holding the [K, OC] window that col2im scatter-adds at t * stride + ggml_tensor * col; + if (depthwise) { + // one group per channel: a batched matmul over the channels scales the kernel by each step + ggml_tensor * krn = ggml_reshape_3d(ctx0, w, 1, K, OC); // [1, K, OC] + ggml_tensor * xs = ggml_reshape_3d(ctx0, x, 1, T, OC); // [1, T, OC] + col = ggml_mul_mat(ctx0, krn, xs); // [K, T, OC] + col = ggml_cont(ctx0, ggml_permute(ctx0, col, 0, 2, 1, 3)); // [K, OC, T] + col = ggml_reshape_2d(ctx0, col, K * OC, T); + } else { + ggml_tensor * w2 = ggml_reshape_2d(ctx0, w, K * OC, w->ne[2]); + w2 = ggml_cont(ctx0, ggml_transpose(ctx0, w2)); // [IC, K * OC] + ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [IC, T] + col = ggml_mul_mat(ctx0, w2, xt); + } + ggml_tensor * full = ggml_col2im_1d(ctx0, col, stride, OC, 0); // [emit_len + p_total, OC] + + ggml_tensor * out; + if (state_name.empty() || p_total == 0) { + out = ggml_cont(ctx0, ggml_view_2d(ctx0, full, emit_len, full->ne[1], full->nb[1], 0)); + } else { + // overlap-add the tail the previous call held back + ggml_tensor * prev = state_in.at(state_name); // [p_total, OC] + ggml_tensor * head = ggml_add(ctx0, ggml_view_2d(ctx0, full, p_total, full->ne[1], full->nb[1], 0), prev); + if (emit_len > p_total) { + ggml_tensor * rest = ggml_view_2d(ctx0, full, emit_len - p_total, full->ne[1], full->nb[1], + (size_t) p_total * full->nb[0]); + out = ggml_concat(ctx0, head, rest, 0); + } else { + out = head; + } + state_out.push_back({state_name, + ggml_cont(ctx0, ggml_view_2d(ctx0, full, p_total, full->ne[1], full->nb[1], + (size_t) emit_len * full->nb[0]))}); + } + + if (b) { + out = ggml_add(ctx0, out, ggml_reshape_2d(ctx0, b, 1, b->ne[0])); + } + return out; +} + +ggml_tensor * clip_graph_pockettts_seanet::res_unit(ggml_tensor * x, const clip_seanet::stage & stage, int dilation, + const std::string & state_prefix) const { + ggml_tensor * h = ggml_elu(ctx0, x); + h = conv1d(h, stage.res_conv1_w, stage.res_conv1_b, 1, dilation, false, state_prefix); + h = ggml_elu(ctx0, h); + // the second conv is pointwise, it needs no left context + h = conv1d(h, stage.res_conv2_w, stage.res_conv2_b, 1, 1); + return ggml_add(ctx0, x, h); +} + +ggml_tensor * clip_graph_pockettts_seanet::encode(ggml_tensor * x) const { + const auto & seanet = model.seanet; + + ggml_tensor * cur = conv1d(x, seanet.conv_in_w, seanet.conv_in_b, 1, 1); + cb(cur, "seanet_enc_in", -1); + + for (int i = 0; i < hparams.seanet_n_stage; i++) { + const auto & stage = seanet.stages[i]; + const int stride = hparams.seanet_ratios[i]; + + cur = res_unit(cur, stage, 1); + cur = ggml_elu(ctx0, cur); + cur = conv1d(cur, stage.scale_conv_w, stage.scale_conv_b, stride, 1); + cb(cur, "seanet_enc_stage", i); + } + + cur = ggml_elu(ctx0, cur); + cur = conv1d(cur, seanet.conv_out_w, seanet.conv_out_b, 1, 1); + cb(cur, "seanet_enc_out", -1); + + return cur; +} + +ggml_tensor * clip_graph_pockettts_seanet::decode(ggml_tensor * x) const { + const auto & seanet = model.seanet; + const bool stream = !state_in.empty(); + + ggml_tensor * cur = conv1d(x, seanet.conv_in_w, seanet.conv_in_b, 1, 1, false, + stream ? "dec_in" : ""); + cb(cur, "seanet_dec_in", -1); + + for (int i = 0; i < hparams.seanet_n_stage; i++) { + const auto & stage = seanet.stages[i]; + // the decoder mirrors the encoder, so the ratios are walked backwards + const int stride = hparams.seanet_ratios[hparams.seanet_n_stage - 1 - i]; + const std::string id = std::to_string(i); + + cur = ggml_elu(ctx0, cur); + cur = conv_transpose1d(cur, stage.scale_conv_w, stage.scale_conv_b, stride, + stream ? "dec_up_" + id : ""); + cur = res_unit(cur, stage, 1, stream ? "dec_res_" + id : ""); + cb(cur, "seanet_dec_stage", i); + } + + cur = ggml_elu(ctx0, cur); + cur = conv1d(cur, seanet.conv_out_w, seanet.conv_out_b, 1, 1, false, + stream ? "dec_out" : ""); + cb(cur, "seanet_dec_out", -1); + + return cur; +} diff --git a/tools/mtmd/models/pockettts-spkenc.cpp b/tools/mtmd/models/pockettts-spkenc.cpp new file mode 100644 index 000000000000..f802d90687d7 --- /dev/null +++ b/tools/mtmd/models/pockettts-spkenc.cpp @@ -0,0 +1,77 @@ +#include "models.h" + +// voice-prompt encoder: raw 24kHz waveform -> one conditioning row per 12.5Hz frame +// mimi encoder (SEANet + transformer + downsample), then flow_lm.speaker_proj_weight + +// pre-norm block with layer scale on both residual paths, see mimi_transformer.py +ggml_tensor * clip_graph_pockettts_spkenc::tfm_layer_forward(ggml_tensor * cur, const clip_layer & layer, ggml_tensor * inp_pos, ggml_tensor * kq_mask, int il) const { + ggml_tensor * inp = cur; + + cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, eps, il); + + ggml_tensor * Qcur = build_mm(layer.q_w, cur); + ggml_tensor * Kcur = build_mm(layer.k_w, cur); + ggml_tensor * Vcur = build_mm(layer.v_w, cur); + + const int64_t n_pos = cur->ne[1]; + Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos); + Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_pos); + Vcur = ggml_reshape_3d(ctx0, Vcur, d_head, n_head, n_pos); + + Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0, + hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0, + hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + + cur = build_attn(layer.o_w, nullptr, Qcur, Kcur, Vcur, kq_mask, kq_scale, il); + cur = ggml_mul(ctx0, cur, layer.ls_1_w); + cur = ggml_add(ctx0, cur, inp); + + inp = cur; + cur = build_norm(cur, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, eps, il); + cur = build_ffn(cur, layer.ff_up_w, nullptr, nullptr, nullptr, layer.ff_down_w, nullptr, FFN_GELU, il); + cur = ggml_mul(ctx0, cur, layer.ls_2_w); + cur = ggml_add(ctx0, cur, inp); + + return cur; +} + +ggml_cgraph * clip_graph_pockettts_spkenc::build() { + // the preprocessor hands over the waveform as a single-row "mel", already [n_samples, 1] + ggml_tensor * inp_raw = build_inp_raw(1); + ggml_tensor * cur = ggml_reshape_2d(ctx0, inp_raw, inp_raw->ne[0], inp_raw->ne[1]); + + clip_graph_pockettts_seanet seanet(*this); + cur = seanet.encode(cur); + cb(cur, "mimi_enc", -1); + + // [T, 512] -> transformer works on [512, T] + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + + ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, cur->ne[1]); + ggml_set_name(inp_pos, "inp_pos"); + ggml_set_input(inp_pos); + + // the mimi transformer is causal with a sliding window, see _build_attention_mask() + ggml_tensor * kq_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, cur->ne[1], cur->ne[1]); + ggml_set_name(kq_mask, "kq_mask"); + ggml_set_input(kq_mask); + + for (int il = 0; il < n_layer; il++) { + cur = tfm_layer_forward(cur, model.layers[il], inp_pos, kq_mask, il); + } + cb(cur, "mimi_enc_tfm", -1); + + // downsample to the model frame rate, [512, T] -> [T, 512] -> [T / 16, 32] + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + cur = seanet.conv1d(cur, model.downsample_w, nullptr, hparams.mimi_downsample, 1, true); + cb(cur, "mimi_downsample", -1); + + // voice latent -> backbone embd + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + cur = build_mm(model.spk_proj_w, cur); + cb(cur, "spk_proj", -1); + + ggml_build_forward_expand(gf, cur); + return gf; +} diff --git a/tools/mtmd/models/qwen3tts-gen.cpp b/tools/mtmd/models/qwen3tts-gen.cpp index b6c95efa941e..84c77f4fad19 100644 --- a/tools/mtmd/models/qwen3tts-gen.cpp +++ b/tools/mtmd/models/qwen3tts-gen.cpp @@ -610,6 +610,10 @@ std::vector list_c2w_state_slots(const clip_hparams & hparams, c const auto & c2w = model.c2w; std::vector slots; + if (c2w.pre_conv_w == nullptr) { + return slots; // not a code2wav model, it keeps no state between calls + } + slots.push_back({"tfm_pos", 1, 1}); // prefix is (W-1) frames, the batch itself gives the other N=W frames (see tfm_layer_forward) diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp index ca4b64efa4a4..98a8c11ee915 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -1423,3 +1423,41 @@ std::vector mtmd_audio_streaming_istft::flush() { return output; } + +// +// mtmd_audio_preprocessor_pockettts +// +// mimi takes the raw 24kHz waveform, there is no mel front-end +// the samples are handed over as a single-row "mel", to reuse the normal chunk path +// + +bool mtmd_audio_preprocessor_pockettts::preprocess(const float * samples, + size_t n_samples, + std::vector & output) { + // the encoder needs whole frames, see pad_for_conv1d() in the reference + const int64_t frame_size = (int64_t) hparams.mimi_downsample * 120; + if (n_samples == 0 || frame_size <= 0) { + return false; + } + + // the mimi transformer mask is dense, so cost is quadratic in the reference length + const int64_t max_samples = (int64_t) clip_hparams::pockettts_max_spk_seconds * hparams.audio_sample_rate; + if ((int64_t) n_samples > max_samples) { + LOG_WRN("%s: speaker reference is %.1f s, truncating to the first %d s\n", __func__, + (double) n_samples / hparams.audio_sample_rate, clip_hparams::pockettts_max_spk_seconds); + n_samples = (size_t) max_samples; + } + + const int64_t n_frames = (int64_t) (n_samples + frame_size - 1) / frame_size; + const int64_t n_padded = n_frames * frame_size; + + mtmd_audio_mel out; + out.n_mel = 1; + out.n_len = n_padded; + out.n_len_org = (int64_t) n_samples; + out.data.assign((size_t) n_padded, 0.0f); + std::copy(samples, samples + n_samples, out.data.begin()); + + output.push_back(std::move(out)); + return true; +} diff --git a/tools/mtmd/mtmd-audio.h b/tools/mtmd/mtmd-audio.h index b4d6f7259808..44ad098ae63a 100644 --- a/tools/mtmd/mtmd-audio.h +++ b/tools/mtmd/mtmd-audio.h @@ -129,6 +129,13 @@ struct mtmd_audio_preprocessor_qwen3tts_spk : mtmd_audio_preprocessor { mtmd_audio_cache cache; }; +// mimi convolves the waveform directly, so this only pads it to a whole number of frames +struct mtmd_audio_preprocessor_pockettts : mtmd_audio_preprocessor { + mtmd_audio_preprocessor_pockettts(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {} + void initialize() override {} + bool preprocess(const float * samples, size_t n_samples, std::vector & output) override; +}; + struct mtmd_audio_preprocessor_parakeet : mtmd_audio_preprocessor { mtmd_audio_preprocessor_parakeet(clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) { } void initialize() override; diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index 85671d1a331f..1c58d3ae1959 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -5,6 +5,8 @@ #include "../src/llama-ext.h" #include +#include +#include #include #include #include @@ -87,7 +89,8 @@ class mtmd_gen_audio_pipeline { virtual int32_t step_prompt(int32_t n_batch) = 0; // sampled can be LLAMA_TOKEN_NULL for pipelines with no discrete backbone token, // those read what they need from h_state_in instead - virtual int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out) = 0; + // set out_stop on end-of-speech, h_state_out must be null if no frame is generated + virtual int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) = 0; virtual int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) = 0; protected: @@ -200,8 +203,10 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { prompt_pos = 0; pos = 0; - top_k = inp->top_k > 0 ? inp->top_k : 50; - top_p = inp->top_p > 0 ? inp->top_p : 1.0f; + const mtmd_gen_inp def = mtmd_gen_inp_default(mctx); + top_k = inp->top_k > 0 ? inp->top_k : def.top_k; + top_p = inp->top_p > 0 ? inp->top_p : def.top_p; + seed = inp->seed; out_type = inp->out_type; // the prompt above holds the whole text stream up to tts_eos, so every generated @@ -241,13 +246,26 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { return n_prompt - prompt_pos; } - int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out) override { - mtmd_gen_inp inp{}; + int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) override { + if (sampled == LLAMA_TOKEN_NULL) { + LOG_ERR("mtmd_helper_gen_audio: qwen3tts requires a token sampled from the backbone\n"); + return 1; + } + + // backbone signals end-of-speech with a token, no frame for this step + if (sampled == codec_eos || llama_vocab_is_eog(vocab, sampled)) { + *out_stop = true; + *h_state_out = nullptr; + return 0; + } + + mtmd_gen_inp inp = mtmd_gen_inp_default(mctx); inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE; inp.code0 = sampled - codec_0; inp.embd = const_cast(h_state_in); inp.top_k = top_k; inp.top_p = top_p; + inp.seed = seed; mtmd_gen_out out{}; if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) { LOG_ERR("mtmd_helper_gen_audio: gen_code process failed\n"); @@ -384,10 +402,11 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { if (codes_buf.empty()) { return true; } - mtmd_gen_inp inp{}; + mtmd_gen_inp inp = mtmd_gen_inp_default(mctx); inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV; inp.codes = codes_buf.data(); inp.n_codes = codes_buf.size(); + inp.seed = seed; // same seed as gen_code, else clip reseeds mid-generation inp.state_data = c2w_state.empty() ? nullptr : (const char *) c2w_state.data(); inp.state_size = c2w_state.size(); mtmd_gen_out out{}; @@ -427,8 +446,9 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { std::unique_ptr prompt_batch; int n_prompt = 0; int prompt_pos = 0; - int32_t top_k = 50; - float top_p = 1.0f; + int32_t top_k = 50; + float top_p = 1.0f; + uint32_t seed = UINT32_MAX; std::vector codes_buf; std::vector c2w_state; std::vector audio_pcm; @@ -438,10 +458,547 @@ class qwen3tts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { std::vector out_buf; }; +// settings that only live in the reference's per-pack yaml, not in the checkpoint +// the english packs share the same shapes and tokenizer, but disagree on these +// all three are 0 / false when the pack does not tune them, the model default is then used +struct pockettts_pack_settings { + float temp = 0.0f; + int frames_after_eos = 0; + bool pad_short_text = false; +}; + +static pockettts_pack_settings pockettts_pack(const char * variant) { + static const std::unordered_map packs = { + { "english", { 0.3f, 0, false } }, + { "english_2026-01", { 0.7f, 0, true } }, + { "english_2026-04", { 0.3f, 0, false } }, + { "french_24l", { 0.7f, 8, false } }, + }; + auto it = packs.find(variant ? variant : ""); + if (it == packs.end()) { + LOG_WRN("mtmd_helper_gen_audio: no tuned settings for pocket-tts variant \"%s\"\n", + variant ? variant : ""); + return {}; + } + return it->second; +} + +// pocket-tts: the backbone emits no token, the flow net turns each hidden state into a latent +// the end-of-speech head also lives in the mmproj +class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { +public: + using mtmd_gen_audio_pipeline::mtmd_gen_audio_pipeline; + + void reset() override { + seq_id = 0; + pos = 0; + feats_buf.clear(); + dec_state.clear(); + audio_pcm.clear(); + h_state_buf.clear(); + out_buf.clear(); + prompt_embd_buf.clear(); + prompt_batch.reset(); + n_prompt = 0; + prompt_pos = 0; + step_idx = 0; + eos_step = -1; + chunks.clear(); + chunk_idx = 0; + n_voice_pos = 0; + chunk_budget = 0; + } + + int32_t set_input(const mtmd_helper_gen_audio_inp * inp) override { + reset(); + seq_id = inp->seq_id; + + if (!ensure_cache()) { + return 1; + } + + std::vector voice; + if (inp->speaker_ref) { + if (!encode_speaker(inp->speaker_ref, voice)) { + return 1; + } + } + + pack = pockettts_pack(info.model_variant); + + const std::string text = prepare_text(std::string(inp->prompt, inp->prompt_len), + pack.pad_short_text); + if (text.empty()) { + LOG_ERR("mtmd_helper_gen_audio: empty prompt\n"); + return 1; + } + + std::vector ids(text.size() + 16); + int n_ids = llama_tokenize(vocab, text.c_str(), (int32_t) text.size(), ids.data(), + (int32_t) ids.size(), false, false); + if (n_ids <= 0) { + LOG_ERR("mtmd_helper_gen_audio: tokenization failed\n"); + return 1; + } + ids.resize((size_t) n_ids); + + // long inputs degrade badly, so each chunk restarts from the voice conditioning + // see split_into_best_sentences() in the reference + chunks = split_chunks(ids); + chunk_idx = 0; + if (chunks.size() > 1) { + LOG_INF("mtmd_helper_gen_audio: %d tokens split into %zu chunks\n", n_ids, chunks.size()); + } + + const int n_e = n_embd; + + // sequence order is voice, then text, then the audio BOS that starts generation + if (!voice.empty()) { + GGML_ASSERT(voice.size() % (size_t) n_e == 0); + if (bos_before_voice != LLAMA_TOKEN_NULL) { + push_embd_row(prompt_embd_buf, bos_before_voice); + } + prompt_embd_buf.insert(prompt_embd_buf.end(), voice.begin(), voice.end()); + } + // every later chunk rewinds to here and re-prompts, so the voice stays primed + n_voice_pos = (int) (prompt_embd_buf.size() / (size_t) n_e); + + for (llama_token t : chunks[0]) { + push_embd_row(prompt_embd_buf, t); + } + push_embd_row(prompt_embd_buf, audio_bos); + arm_chunk_budget(0); + + n_prompt = (int) (prompt_embd_buf.size() / (size_t) n_e); + prompt_batch.reset(new decode_embd_batch(prompt_embd_buf.data(), n_prompt, 1, n_e)); + prompt_batch->set_position_normal(0, seq_id); + prompt_pos = 0; + + seed = inp->seed; + out_type = inp->out_type; + + return 0; + } + + int32_t step_prompt(int32_t n_batch) override { + GGML_ASSERT(n_batch > 0); + if (prompt_pos >= n_prompt) { + return 0; + } + const int32_t n_tokens_batch = std::min(n_batch, n_prompt - prompt_pos); + llama_batch batch_view = prompt_batch->get_view(prompt_pos, n_tokens_batch); + + if ((prompt_pos + n_tokens_batch) == n_prompt) { + batch_view.logits[n_tokens_batch - 1] = 1; + } + + if (llama_decode(lctx, batch_view) != 0) { + LOG_ERR("mtmd_helper_gen_audio: prompt decode failed\n"); + return -1; + } + + pos += n_tokens_batch; + prompt_pos += n_tokens_batch; + + if (prompt_pos >= n_prompt) { + prompt_batch.reset(); + prompt_embd_buf.clear(); + return 0; + } + return n_prompt - prompt_pos; + } + + int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) override { + (void) sampled; // the backbone output is continuous, there is no token to consume + + mtmd_gen_inp inp = mtmd_gen_inp_default(mctx); + inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE; + inp.embd = const_cast(h_state_in); + // clip only reseeds when the seed changes, so pass the same one on every step + inp.seed = seed; + if (pack.temp > 0.0f) { + inp.temp = pack.temp; + } + mtmd_gen_out out{}; + if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) { + LOG_ERR("mtmd_helper_gen_audio: flow decode failed\n"); + return 1; + } + if (out.is_eos && eos_step < 0) { + eos_step = step_idx; + } + // the frame of the stopping step is discarded, matching _autoregressive_generation(). + // the budget is the reference's fallback for a chunk whose eos head never fires + const bool chunk_done = (eos_step >= 0 && step_idx >= eos_step + frames_after_eos) || + step_idx >= chunk_budget; + if (chunk_done) { + if (eos_step < 0) { + LOG_WRN("mtmd_helper_gen_audio: chunk %zu hit its budget without end-of-speech\n", chunk_idx); + } + return finish_chunk(h_state_out, out_stop); + } + + feats_buf.insert(feats_buf.end(), out.feats, out.feats + out.n_feats); + step_idx++; + if (out.n_feats > 0 && feats_buf.size() / out.n_feats >= window_frames) { + if (!flush_gen_wav()) { + return 1; + } + } + + decode_embd_batch batch_embd(const_cast(out.embd), 1, 1, n_embd); + batch_embd.set_position_normal(pos, seq_id); + batch_embd.batch.logits[0] = 1; + pos++; + + if (llama_decode(lctx, batch_embd.batch) != 0) { + LOG_ERR("mtmd_helper_gen_audio: decode failed\n"); + return 1; + } + + const float * he = llama_get_embeddings_ith(lctx, -1); + h_state_buf.assign(he, he + n_embd); + *h_state_out = h_state_buf.data(); + + return 0; + } + + int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) override { + if (!flush_gen_wav()) { + return 1; + } + + *out_sample_rate = info.sample_rate; + if (out_n_samples) { + *out_n_samples = (int64_t) audio_pcm.size(); + } + + if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) { + *out_data = (const char *) audio_pcm.data(); + *out_data_len = audio_pcm.size() * sizeof(float); + return 0; + } + + out_buf.clear(); + if (!write_wav16(out_buf, audio_pcm, info.sample_rate)) { + LOG_ERR("mtmd_helper_gen_audio: output too large for WAV\n"); + return 1; + } + *out_data = out_buf.data(); + *out_data_len = out_buf.size(); + return 0; + } + +private: + bool ensure_cache() { + if (specials_ok) { + return true; + } + // bos_before_voice is optional, some packs do not insert it + bos_before_voice = find_special_token(vocab, "<|bos_before_voice|>"); + audio_bos = find_special_token(vocab, "<|audio_bos|>"); + if (audio_bos == LLAMA_TOKEN_NULL) { + LOG_ERR("mtmd_helper_gen_audio: missing <|audio_bos|> in vocab\n"); + return false; + } + const uint32_t n_tok_embd = llama_model_get_tok_embd(model, nullptr); + if (n_tok_embd == 0) { + LOG_ERR("mtmd_helper_gen_audio: model has no token embeddings\n"); + return false; + } + tok_embd.resize(n_tok_embd); + if (llama_model_get_tok_embd(model, tok_embd.data()) != n_tok_embd) { + LOG_ERR("mtmd_helper_gen_audio: token embedding copy failed\n"); + return false; + } + GGML_ASSERT(n_embd > 0 && n_tok_embd % (uint32_t) n_embd == 0); + specials_ok = true; + return true; + } + + // the table can be shorter than the vocab, so bound the row lookup + void push_embd_row(std::vector & dst, llama_token t) const { + const size_t n_rows = tok_embd.size() / (size_t) n_embd; + GGML_ASSERT(t >= 0 && (size_t) t < n_rows); + dst.insert(dst.end(), + tok_embd.begin() + (size_t) t * n_embd, + tok_embd.begin() + (size_t) (t + 1) * n_embd); + } + + // token ids of the pieces the reference splits on, see split_into_best_sentences(). + // the leading token is dropped, it is the tokenizer's dummy prefix + std::vector punct_ids(const char * s) const { + std::vector ids(16); + const int n = llama_tokenize(vocab, s, (int32_t) strlen(s), ids.data(), (int32_t) ids.size(), false, false); + if (n <= 1) { + return {}; + } + return std::vector(ids.begin() + 1, ids.begin() + n); + } + + // cut after runs of boundary tokens, so punctuation stays with the sentence it ends + static std::vector> split_on(const std::vector & ids, + const std::vector & boundary) { + std::vector> out; + size_t start = 0; + bool prev_was_boundary = false; + for (size_t i = 0; i < ids.size(); i++) { + const bool is_boundary = std::find(boundary.begin(), boundary.end(), ids[i]) != boundary.end(); + if (!is_boundary && prev_was_boundary) { + out.emplace_back(ids.begin() + start, ids.begin() + i); + start = i; + } + prev_was_boundary = is_boundary; + } + out.emplace_back(ids.begin() + start, ids.end()); + return out; + } + + std::vector> split_chunks(const std::vector & ids) const { + if ((int) ids.size() <= max_chunk_tokens) { + return { ids }; + } + const std::vector eos_punct = punct_ids(".!...?"); + const std::vector mid_punct = punct_ids(",;:"); + + // oversized sentences are split again on weaker punctuation, else words get skipped + std::vector> segments; + for (auto & seg : split_on(ids, eos_punct)) { + if ((int) seg.size() <= max_chunk_tokens) { + segments.push_back(std::move(seg)); + continue; + } + auto sub = split_on(seg, mid_punct); + if (sub.size() > 1) { + for (auto & s : sub) { + segments.push_back(std::move(s)); + } + } else { + segments.push_back(std::move(seg)); + } + } + + std::vector> out; + for (auto & seg : segments) { + if (seg.empty()) { + continue; + } + if (!out.empty() && (int) (out.back().size() + seg.size()) <= max_chunk_tokens) { + out.back().insert(out.back().end(), seg.begin(), seg.end()); + } else { + out.push_back(std::move(seg)); + } + } + if (out.empty()) { + out.push_back(ids); + } + for (const auto & c : out) { + if ((int) c.size() > max_chunk_tokens) { + LOG_WRN("mtmd_helper_gen_audio: chunk of %zu tokens exceeds the %d token budget, " + "generation may skip words\n", c.size(), max_chunk_tokens); + } + } + return out; + } + + // _estimate_max_gen_len() plus the per-chunk tail guess, both in frames + void arm_chunk_budget(size_t idx) { + const int n_tok = (int) chunks[idx].size(); + chunk_budget = (int) std::ceil((n_tok / 3.0 + 2.0) * frame_rate); + // the pack may pin the tail, else the reference guesses it from the word count + frames_after_eos = pack.frames_after_eos > 0 ? pack.frames_after_eos : (n_tok <= 6 ? 5 : 3); + step_idx = 0; + eos_step = -1; + } + + // ends the current chunk and, if there is another, re-prompts it on top of the voice + int32_t finish_chunk(const float ** h_state_out, bool * out_stop) { + if (!flush_gen_wav()) { + return 1; + } + // the decoder restarts too, the next chunk's audio is not continuous with this one + dec_state.clear(); + + if (chunk_idx + 1 >= chunks.size()) { + *out_stop = true; + *h_state_out = nullptr; + return 0; + } + chunk_idx++; + + // drop this chunk's text and audio, keep the voice conditioning + llama_memory_seq_rm(llama_get_memory(lctx), seq_id, n_voice_pos, -1); + pos = n_voice_pos; + + const int n_e = n_embd; + prompt_embd_buf.clear(); + for (llama_token t : chunks[chunk_idx]) { + push_embd_row(prompt_embd_buf, t); + } + push_embd_row(prompt_embd_buf, audio_bos); + arm_chunk_budget(chunk_idx); + + const int n_rows = (int) (prompt_embd_buf.size() / (size_t) n_e); + GGML_ASSERT(n_rows > 0); + decode_embd_batch batch(prompt_embd_buf.data(), n_rows, 1, n_e); + batch.set_position_normal(pos, seq_id); + batch.batch.logits[n_rows - 1] = 1; + if (llama_decode(lctx, batch.batch) != 0) { + LOG_ERR("mtmd_helper_gen_audio: chunk prompt decode failed\n"); + return 1; + } + pos += n_rows; + prompt_embd_buf.clear(); + + const float * he = llama_get_embeddings_ith(lctx, -1); + h_state_buf.assign(he, he + n_embd); + *h_state_out = h_state_buf.data(); + *out_stop = false; + return 0; + } + + // same normalization as prepare_text_prompt() in the reference, it affects quality + static std::string prepare_text(const std::string & in, bool pad_short) { + std::string s; + s.reserve(in.size() + 1); + for (char c : in) { + if (c == '\n' || c == '\r') { + s += ' '; + } else if (c == ';') { + s += ','; + } else { + s += c; + } + } + const size_t b = s.find_first_not_of(' '); + const size_t e = s.find_last_not_of(' '); + if (b == std::string::npos) { + return ""; + } + s = s.substr(b, e - b + 1); + if (s[0] >= 'a' && s[0] <= 'z') { + s[0] = (char) (s[0] - 'a' + 'A'); + } + const unsigned char last = (unsigned char) s.back(); + if (std::isalnum(last)) { + s += '.'; + } + if (pad_short && count_words(s) < 5) { + s = std::string(8, ' ') + s; + } + return s; + } + + static int count_words(const std::string & s) { + int n = 0; + bool in_word = false; + for (char c : s) { + if (c == ' ') { + in_word = false; + } else if (!in_word) { + in_word = true; + n++; + } + } + return n; + } + + // runs the reference wav through the mimi encoder, returns one row per 12.5Hz frame + bool encode_speaker(mtmd_bitmap * bitmap, std::vector & out) { + if (!mtmd_support_audio(mctx)) { + LOG_ERR("mtmd_helper_gen_audio: mmproj has no voice encoder\n"); + return false; + } + const std::string marker = mtmd_default_marker(); + mtmd_input_text text{ marker.c_str(), marker.size(), false, true }; + mtmd_input_chunks * chunks = mtmd_input_chunks_init(); + const mtmd_bitmap * bptr = bitmap; + bool ok = mtmd_tokenize(mctx, chunks, &text, &bptr, 1) == 0; + if (ok) { + ok = false; + for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) { + const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i); + if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_AUDIO) { + continue; + } + if (mtmd_encode_chunk(mctx, chunk) != 0) { + LOG_ERR("mtmd_helper_gen_audio: voice encode failed\n"); + break; + } + const float * embd = mtmd_get_output_embd(mctx); + const size_t n = (size_t) llama_model_n_embd_inp(model) * mtmd_input_chunk_get_n_tokens(chunk); + out.assign(embd, embd + n); + ok = true; + break; + } + } + mtmd_input_chunks_free(chunks); + return ok; + } + + // decodes the buffered latents, the mimi decoder state carries over between calls + bool flush_gen_wav() { + if (feats_buf.empty()) { + return true; + } + mtmd_gen_inp inp = mtmd_gen_inp_default(mctx); + inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV; + inp.feats = feats_buf.data(); + inp.n_feats = feats_buf.size(); + inp.seed = seed; + inp.state_data = dec_state.empty() ? nullptr : (const char *) dec_state.data(); + inp.state_size = dec_state.size(); + mtmd_gen_out out{}; + if (mtmd_gen_audio_process(mctx, &inp, &out) != 0) { + LOG_ERR("mtmd_helper_gen_audio: mimi decode failed\n"); + return false; + } + audio_pcm.insert(audio_pcm.end(), out.audio, out.audio + out.n_samples); + dec_state.assign(out.state_data, out.state_data + out.state_size); + feats_buf.clear(); + return true; + } + + pockettts_pack_settings pack; + bool specials_ok = false; + llama_token bos_before_voice = LLAMA_TOKEN_NULL; + llama_token audio_bos = LLAMA_TOKEN_NULL; + std::vector tok_embd; + + llama_seq_id seq_id = 0; + int pos = 0; + std::vector prompt_embd_buf; + std::unique_ptr prompt_batch; + int n_prompt = 0; + int prompt_pos = 0; + uint32_t seed = UINT32_MAX; + // end-of-speech is latched, then a few more frames are generated as tail padding + int step_idx = 0; + int eos_step = -1; + int frames_after_eos = 3; + static constexpr int max_chunk_tokens = 50; // MAX_TOKEN_PER_CHUNK in the reference + static constexpr double frame_rate = 12.5; + std::vector> chunks; + size_t chunk_idx = 0; + int n_voice_pos = 0; // KV positions held by the voice conditioning + int chunk_budget = 0; + + // latents are decoded a window at a time, the decoder state bridges the windows + size_t window_frames = 8; + std::vector feats_buf; + std::vector dec_state; + std::vector audio_pcm; + std::vector h_state_buf; + mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV; + std::vector out_buf; +}; + static std::unique_ptr make_pipeline(llama_context * lctx, mtmd_context * mctx) { switch (mtmd_gen_audio_get_info(mctx).type) { case MTMD_GEN_AUDIO_TYPE_QWEN3TTS: return std::unique_ptr(new qwen3tts_gen_audio_pipeline(lctx, mctx)); + case MTMD_GEN_AUDIO_TYPE_POCKETTTS: + return std::unique_ptr(new pockettts_gen_audio_pipeline(lctx, mctx)); default: return nullptr; } @@ -483,11 +1040,17 @@ int32_t mtmd_helper_gen_audio_step_prompt(mtmd_helper_gen_audio * ctx, int32_t n } int32_t mtmd_helper_gen_audio_step_gen(mtmd_helper_gen_audio * ctx, llama_token sampled, - const float * h_state_in, const float ** h_state_out) { + const float * h_state_in, const float ** h_state_out, + bool * out_stop) { if (!ctx->pipeline) { return 1; } - return ctx->pipeline->step_gen(sampled, h_state_in, h_state_out); + bool stop = false; + const int32_t ret = ctx->pipeline->step_gen(sampled, h_state_in, h_state_out, &stop); + if (out_stop) { + *out_stop = stop; + } + return ret; } int32_t mtmd_helper_gen_audio_get_output(mtmd_helper_gen_audio * ctx, int32_t * out_sample_rate, diff --git a/tools/mtmd/mtmd-helper.h b/tools/mtmd/mtmd-helper.h index 7e5cf9b5098c..832f7171ac71 100644 --- a/tools/mtmd/mtmd-helper.h +++ b/tools/mtmd/mtmd-helper.h @@ -183,8 +183,9 @@ struct mtmd_helper_gen_audio_inp { mtmd_bitmap * speaker_ref; // optional, can be NULL const char * lang; // optional, can be NULL - int32_t top_k; - float top_p; + int32_t top_k; + float top_p; + uint32_t seed; // UINT32_MAX for random (default: random) enum mtmd_helper_gen_audio_outtype out_type; }; @@ -208,12 +209,15 @@ MTMD_API int32_t mtmd_helper_gen_audio_step_prompt( int32_t n_batch); // generates one frame; must only be called after step_prompt() has returned 0 -// h_state_out is valid until next step_gen() or reset() call +// sampled can be LLAMA_TOKEN_NULL for pipelines with no discrete backbone token +// out_stop (optional) is set on end-of-speech, the caller must then stop the loop +// h_state_out is valid until next step_gen() or reset() call, null if no frame is generated MTMD_API int32_t mtmd_helper_gen_audio_step_gen( mtmd_helper_gen_audio * ctx, llama_token sampled, const float * h_state_in, - const float ** h_state_out); + const float ** h_state_out, + bool * out_stop); // out_data valid until next get_output() or reset() call // out_n_samples (optional, can be NULL) receives the number of generated PCM samples @@ -261,8 +265,8 @@ struct gen_audio { int32_t step_prompt(int32_t n_batch) { return mtmd_helper_gen_audio_step_prompt(ctx.get(), n_batch); } - int32_t step_gen(llama_token sampled, const float * h_state, const float ** h_state_out) { - return mtmd_helper_gen_audio_step_gen(ctx.get(), sampled, h_state, h_state_out); + int32_t step_gen(llama_token sampled, const float * h_state, const float ** h_state_out, bool * out_stop = nullptr) { + return mtmd_helper_gen_audio_step_gen(ctx.get(), sampled, h_state, h_state_out, out_stop); } int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples = nullptr) { return mtmd_helper_gen_audio_get_output(ctx.get(), out_sample_rate, out_data, out_data_len, out_n_samples); diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 82b73d5cd126..4b9c45d62677 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -477,6 +477,7 @@ struct mtmd_context { // generation context struct clip_ctx * ctx_gen_a; // audio std::vector gen_out_codes; // this frame's 16 sampled codes (GEN_CODE) + std::vector gen_out_feats; // this frame's continuous features, if any (GEN_CODE) std::vector gen_out_embd; // next-step hidden state fed back to backbone (GEN_CODE) std::vector gen_out_audio; // decoded PCM samples for the current frame (GEN_WAV) std::vector gen_out_state; // state to feed into the next GEN_WAV call @@ -979,6 +980,10 @@ struct mtmd_context { { audio_preproc = std::make_unique(ctx_a); } break; + case PROJECTOR_TYPE_POCKETTTS_SPKENC: + { + audio_preproc = std::make_unique(ctx_a); + } break; default: throw std::runtime_error(string_format("%s: unexpected audio projector type %d\n", __func__, proj)); } @@ -1798,16 +1803,22 @@ float * mtmd_get_output_embd(mtmd_context * ctx) { // mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) { - mtmd_gen_audio_info info; + mtmd_gen_audio_info info{}; + info.model_variant = ""; if (!ctx->ctx_gen_a) { info.type = MTMD_GEN_AUDIO_TYPE_NONE; return info; } + info.model_variant = clip_get_hparams(ctx->ctx_gen_a)->gen_model_variant.c_str(); switch (clip_get_projector_type(ctx->ctx_gen_a)) { case PROJECTOR_TYPE_QWEN3TTS_GEN: info.type = MTMD_GEN_AUDIO_TYPE_QWEN3TTS; info.sample_rate = 24000; break; + case PROJECTOR_TYPE_POCKETTTS_GEN: + info.type = MTMD_GEN_AUDIO_TYPE_POCKETTTS; + info.sample_rate = 24000; + break; default: info.type = MTMD_GEN_AUDIO_TYPE_NONE; break; @@ -1815,6 +1826,33 @@ mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) { return info; } +mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx) { + mtmd_gen_inp inp{}; + inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE; + inp.seed = UINT32_MAX; + if (!ctx->ctx_gen_a) { + return inp; + } + + switch (clip_get_projector_type(ctx->ctx_gen_a)) { + case PROJECTOR_TYPE_QWEN3TTS_GEN: + // https://huggingface.co/Qwen/Qwen3-TTS-12Hz-1.7B-Base/blob/main/generation_config.json + inp.top_k = 50; + inp.top_p = 1.0f; + inp.temp = 0.9f; // TODO: handle this on graph + break; + case PROJECTOR_TYPE_POCKETTTS_GEN: + // https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/default_parameters.py + inp.top_k = 50; + inp.top_p = 1.0f; + inp.temp = 0.7f; + break; + default: + break; + } + return inp; +} + static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_inp * inp, mtmd_gen_out * out) { clip_ctx * ctx_clip = ctx->ctx_gen_a; if (!ctx_clip) { @@ -1822,6 +1860,8 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in return 1; } + *out = {}; + if (inp->type == MTMD_GEN_PROCESS_TYPE_GEN_CODE) { const size_t n_embd = (size_t) clip_n_mmproj_embd(ctx_clip); @@ -1835,16 +1875,22 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in std::vector out_embd(n_embd); std::vector out_codes; + std::vector out_feats; + bool is_eos = false; clip_encode_params params; - params.imgs = &batch; - params.n_threads = ctx->n_threads; - params.gen_process = CLIP_GEN_PROCESS_GEN_CODE; - params.out_embd = &out_embd; - params.out_codes = &out_codes; - params.code0 = inp->code0; - params.top_k = inp->top_k; - params.top_p = inp->top_p; + params.imgs = &batch; + params.n_threads = ctx->n_threads; + params.gen_process = CLIP_GEN_PROCESS_GEN_CODE; + params.out_embd = &out_embd; + params.out_codes = &out_codes; + params.out_feats = &out_feats; + params.code0 = inp->code0; + params.top_k = inp->top_k; + params.top_p = inp->top_p; + params.seed = inp->seed; + params.temp = inp->temp; + params.out_is_eos = &is_eos; if (!clip_encode(ctx_clip, ¶ms)) { LOG_ERR("%s: clip_encode failed (gen_code)\n", __func__); @@ -1853,19 +1899,31 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in ctx->gen_out_embd = std::move(out_embd); ctx->gen_out_codes = std::move(out_codes); - - out->embd = ctx->gen_out_embd.data(); - out->codes = ctx->gen_out_codes.data(); - out->n_codes = ctx->gen_out_codes.size(); + ctx->gen_out_feats = std::move(out_feats); + + out->embd = ctx->gen_out_embd.data(); + out->codes = ctx->gen_out_codes.data(); + out->n_codes = ctx->gen_out_codes.size(); + out->feats = ctx->gen_out_feats.data(); + out->n_feats = ctx->gen_out_feats.size(); + out->is_eos = is_eos; return 0; } // MTMD_GEN_PROCESS_TYPE_GEN_WAV - if (!inp->codes || inp->n_codes == 0) { - LOG_ERR("%s: codes required for gen_wav\n", __func__); + const bool has_codes = inp->codes && inp->n_codes > 0; + const bool has_feats = inp->feats && inp->n_feats > 0; + if (has_codes == has_feats) { + LOG_ERR("%s: gen_wav requires exactly one of codes or feats\n", __func__); return 1; } - std::vector in_codes(inp->codes, inp->codes + inp->n_codes); + std::vector in_codes; + std::vector in_feats; + if (has_codes) { + in_codes.assign(inp->codes, inp->codes + inp->n_codes); + } else { + in_feats.assign(inp->feats, inp->feats + inp->n_feats); + } std::vector in_state; if (inp->state_data) { in_state.assign(inp->state_data, inp->state_data + inp->state_size); @@ -1885,7 +1943,10 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in params.imgs = &batch; params.n_threads = ctx->n_threads; params.gen_process = CLIP_GEN_PROCESS_GEN_WAV; - params.codes = &in_codes; + // gen_wav draws no randomness, but keep the seed so it does not reseed mid-generation + params.seed = inp->seed; + params.codes = has_codes ? &in_codes : nullptr; + params.feats = has_feats ? &in_feats : nullptr; params.out_audio = &ctx->gen_out_audio; params.state_in = inp->state_data ? &in_state : nullptr; params.state_out = &ctx->gen_out_state; diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index e5063d114cc9..c1a5921db2f3 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -344,18 +344,25 @@ MTMD_API struct mtmd_caps mtmd_get_cap_from_file(const char * mmproj_fname); enum mtmd_gen_audio_type { MTMD_GEN_AUDIO_TYPE_NONE, // not supported MTMD_GEN_AUDIO_TYPE_QWEN3TTS, + MTMD_GEN_AUDIO_TYPE_POCKETTTS, }; + struct mtmd_gen_audio_info { enum mtmd_gen_audio_type type; int32_t sample_rate; // in Hz, for example 24000 for qwen3tts + const char * model_variant; // name of the weight variant, can be nullptr if not applicable }; + MTMD_API struct mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx); + enum mtmd_gen_process_type { MTMD_GEN_PROCESS_TYPE_GEN_CODE, // h_state to semantic (codes, mel-spectrogram, etc.) MTMD_GEN_PROCESS_TYPE_GEN_WAV, // convert semantic to PCM audio // for qwen3tts, this is code2wav + // for pocket-tts, this is mimi decoder }; + struct mtmd_gen_inp { enum mtmd_gen_process_type type; @@ -364,21 +371,30 @@ struct mtmd_gen_inp { float * embd; // the hidden state from backbone, must have n_text_embd elements int32_t top_k; float top_p; + uint32_t seed; // UINT32_MAX for random + float temp; // sampling temperature, or noise scale for flow-matching decoders // for MTMD_GEN_PROCESS_TYPE_GEN_WAV + // pass either codes (discrete) or feats (continuous), depending on the pipeline int32_t * codes; size_t n_codes; + const float * feats; + size_t n_feats; const char * state_data; size_t state_size; }; + struct mtmd_gen_out { // note: output memory is allocated by the context, valid until next process() call // for MTMD_GEN_PROCESS_TYPE_GEN_CODE const int32_t * codes; - size_t n_codes; + size_t n_codes; + const float * feats; // continuous counterpart of codes + size_t n_feats; const float * embd; // the generated hidden state, to be fed back to backbone // it must have n_text_embd elements + bool is_eos; // only set by pipelines having the EOS head inside mmproj // for MTMD_GEN_PROCESS_TYPE_GEN_WAV const float * audio; @@ -386,6 +402,10 @@ struct mtmd_gen_out { const char * state_data; size_t state_size; }; + +// defaults tuned for the loaded pipeline, callers override only what they care about +MTMD_API struct mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx); + // note: this API is stateless, caller must handle state management and audio frame accumulation MTMD_API int32_t mtmd_gen_audio_process(mtmd_context * ctx, const struct mtmd_gen_inp * inp, diff --git a/tools/mtmd/requirements.txt b/tools/mtmd/requirements.txt index f26d8e912a37..d646ca7b02f2 100644 --- a/tools/mtmd/requirements.txt +++ b/tools/mtmd/requirements.txt @@ -2,11 +2,5 @@ --extra-index-url https://download.pytorch.org/whl/cpu pillow~=11.3.0 -## Embedding Gemma requires PyTorch 2.6.0 or later, bumped to 2.11.0 for compatibility -torch==2.11.0; platform_machine != "s390x" # check_requirements: ignore "==" +torch==2.11.0 # check_requirements: ignore "==" torchvision==0.26.0; platform_machine != "s390x" # check_requirements: ignore "==" - -# torch s390x packages can only be found from nightly builds ---extra-index-url https://download.pytorch.org/whl/nightly -torch>=0.0.0.dev0; platform_machine == "s390x" # check_requirements: ignore "==" -torchvision>=0.0.0.dev0; platform_machine == "s390x" # check_requirements: ignore "==" diff --git a/tools/server/README.md b/tools/server/README.md index 6927caddbb6c..a2ab872b4c1a 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -75,7 +75,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing
(env: LLAMA_ARG_MLOCK) | | `--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: mmap)
- 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) | +| `-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) | | `--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 | diff --git a/tools/server/tests/tests.sh b/tools/server/tests/tests.sh index 8d6681193db3..433dc99828e4 100755 --- a/tools/server/tests/tests.sh +++ b/tools/server/tests/tests.sh @@ -6,11 +6,6 @@ cd $SCRIPT_DIR set -eu -if [[ "${SLOW_TESTS:-0}" == 1 ]]; then - # Slow tests for tool calls need quite a few models ahead of time to avoid timing out. - python $SCRIPT_DIR/../../../scripts/fetch_server_test_models.py -fi - if [ $# -lt 1 ] then if [[ "${SLOW_TESTS:-0}" == 1 ]]; then diff --git a/tools/tts/README.md b/tools/tts/README.md index dd84336c3990..1b08d5ef3218 100644 --- a/tools/tts/README.md +++ b/tools/tts/README.md @@ -32,3 +32,28 @@ llama-tts -hf ggml-org/Qwen3-TTS-12Hz-1.7B-Base-GGUF \ --tts-speaker-file speaker.mp3 \ --output out.wav ``` + +## Pocket TTS + +Available params: +- `--tts-speaker-file` should point to a speaker reference audio file (wav, mp3). It is required, the model produces almost no audio without it +- Note: `lang` is not used, the language is a property of the weights + +Example usage: + +```sh +llama-tts -m pocket-tts.gguf \ + -mm mmproj-pocket-tts.gguf \ + -p "Hello world" \ + --tts-speaker-file speaker.mp3 \ + --output out.wav +``` + +**Note for GGUF conversion:** + +The [upstream repository](https://huggingface.co/kyutai/pocket-tts) holds one complete model per language under `languages/`, next to a set of shared files at the root. Convert one of the `languages/` directories, **not** the root directory: + +```sh +python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --outfile pocket-tts.gguf +python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --mmproj --outfile mmproj-pocket-tts.gguf +``` diff --git a/tools/tts/tts.cpp b/tools/tts/tts.cpp index 49c405e14784..fd7522f8defd 100644 --- a/tools/tts/tts.cpp +++ b/tools/tts/tts.cpp @@ -119,6 +119,7 @@ int main(int argc, char ** argv) { inp.lang = params.tts_lang.c_str(); inp.top_k = params.sampling.top_k; inp.top_p = params.sampling.top_p; + inp.seed = params.sampling.seed; inp.out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV; // @@ -143,8 +144,7 @@ int main(int argc, char ** argv) { } } - const llama_vocab * vocab = llama_model_get_vocab(model); - + // note: some pipelines ignore this token and use the hidden state instead auto sample_semantic_code = [&]() -> llama_token { llama_token t = common_sampler_sample(smpl, lctx, -1); common_sampler_accept(smpl, t, true); @@ -159,19 +159,24 @@ int main(int argc, char ** argv) { tts_timings timings; const int64_t t_gen_start_us = ggml_time_us(); - for (; n_frames < max_new && !llama_vocab_is_eog(vocab, sampled); n_frames++) { + bool stop = false; + while (!stop && n_frames < max_new) { const float * h_next = nullptr; // stage 2+3: semantic --> acoustic details --> audio waveform // step_gen() runs both stages and returns new h_state for next step - if (gen.step_gen(sampled, h_state, &h_next) != 0) { + if (gen.step_gen(sampled, h_state, &h_next, &stop) != 0) { LOG_ERR("step_gen failed at frame %d\n", n_frames); return 1; } + if (!h_next) { + break; // stopped without generating a frame + } + n_frames++; h_state = h_next; sampled = sample_semantic_code(); - timings.report(n_frames + 1); + timings.report(n_frames); } const double t_gen_s = (ggml_time_us() - t_gen_start_us) / 1e6; diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetails.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetails.svelte index ab36e4a43f74..eaaba69de622 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetails.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetails.svelte @@ -3,6 +3,7 @@ import { ChevronDown } from '@lucide/svelte'; import * as Collapsible from '$lib/components/ui/collapsible'; import { STATS_UNITS } from '$lib/constants'; + import { gaugePopup } from '$lib/stores/context-gauge-popup.svelte'; interface Props { currentRead: number; @@ -30,19 +31,19 @@ transientDetails }: Props = $props(); - let open = $state(false); - const hasCumulative = $derived(cumulativeRead > 0 || cumulativeOutput > 0); const hasCurrent = $derived(currentRead > 0 || currentOutput > 0); - + Token usage details - + diff --git a/tools/ui/src/lib/hooks/use-context-gauge.svelte.ts b/tools/ui/src/lib/hooks/use-context-gauge.svelte.ts index 76dcd356ec29..7ab3f83df609 100644 --- a/tools/ui/src/lib/hooks/use-context-gauge.svelte.ts +++ b/tools/ui/src/lib/hooks/use-context-gauge.svelte.ts @@ -112,7 +112,7 @@ export function useContextGauge(): UseContextGaugeReturn { return chatStore.getConversationModel(activeMessages() as DatabaseMessage[]); }); const isActiveModelLoaded = $derived( - activeModelId !== null && modelsStore.isModelLoaded(activeModelId) + activeModelId !== null && (!isRouterMode() || modelsStore.isModelLoaded(activeModelId)) ); const isActiveModelLoading = $derived( activeModelId !== null && modelsStore.isModelOperationInProgress(activeModelId) diff --git a/tools/ui/src/lib/stores/context-gauge-popup.svelte.ts b/tools/ui/src/lib/stores/context-gauge-popup.svelte.ts index 654dc9f7787b..34f25283f89e 100644 --- a/tools/ui/src/lib/stores/context-gauge-popup.svelte.ts +++ b/tools/ui/src/lib/stores/context-gauge-popup.svelte.ts @@ -16,7 +16,7 @@ import { let closeTimer: ReturnType | undefined; let lastPointerType = ''; -export const gaugePopup = $state({ bottom: 0, centerX: 0, open: false }); +export const gaugePopup = $state({ bottom: 0, centerX: 0, detailsOpen: false, open: false }); function openFrom(trigger: HTMLElement): void { clearTimeout(closeTimer);